From 4390ebf5a9e7977583b976ff206a7f7dee3f93a0 Mon Sep 17 00:00:00 2001 From: A Date: Sat, 11 Jul 2026 19:48:26 +0800 Subject: [PATCH] perf: sync protocol and core hardening updates --- .env.example | 40 + cmd/telesrv/main.go | 54 +- .../0062_private_message_idempotency.down.sql | 3 + .../0062_private_message_idempotency.up.sql | 25 + .../0063_channel_update_retention.down.sql | 19 + .../0063_channel_update_retention.up.sql | 35 + .../0064_user_update_retention.down.sql | 20 + .../0064_user_update_retention.up.sql | 22 + .../0065_auth_key_orphan_retention.down.sql | 1 + .../0065_auth_key_orphan_retention.up.sql | 4 + .../0066_dispatch_outbox_user_lanes.down.sql | 2 + .../0066_dispatch_outbox_user_lanes.up.sql | 39 + .../0067_auth_key_last_used.down.sql | 7 + .../migrations/0067_auth_key_last_used.up.sql | 17 + ...0068_private_message_send_receipt.down.sql | 5 + .../0068_private_message_send_receipt.up.sql | 15 + .../0069_dispatch_outbox_user_heads.down.sql | 12 + .../0069_dispatch_outbox_user_heads.up.sql | 80 ++ ...70_dispatch_outbox_head_readiness.down.sql | 46 + ...0070_dispatch_outbox_head_readiness.up.sql | 82 ++ .../0071_temp_auth_key_expiry_seek.down.sql | 1 + .../0071_temp_auth_key_expiry_seek.up.sql | 5 + ...ispatch_outbox_head_ready_indexes.down.sql | 16 + ..._dispatch_outbox_head_ready_indexes.up.sql | 35 + .../0073_send_replay_snapshots.down.sql | 17 + .../0073_send_replay_snapshots.up.sql | 23 + ...dispatch_outbox_poison_head_index.down.sql | 1 + ...4_dispatch_outbox_poison_head_index.up.sql | 10 + .../0075_uploaded_media_receipts.down.sql | 1 + .../0075_uploaded_media_receipts.up.sql | 12 + .../0076_album_group_reservations.down.sql | 1 + .../0076_album_group_reservations.up.sql | 17 + ...077_correct_private_send_defaults.down.sql | 8 + .../0077_correct_private_send_defaults.up.sql | 17 + .../0078_channel_send_fingerprint.down.sql | 3 + .../0078_channel_send_fingerprint.up.sql | 16 + ...079_login_code_message_deliveries.down.sql | 1 + .../0079_login_code_message_deliveries.up.sql | 22 + .../0080_login_code_delivery_expiry.down.sql | 4 + .../0080_login_code_delivery_expiry.up.sql | 15 + internal/app/account/login_email_test.go | 283 +++- internal/app/account/phone_change.go | 70 +- internal/app/account/phone_change_test.go | 45 +- internal/app/account/service.go | 167 ++- internal/app/auth/change_phone_resend_test.go | 32 + internal/app/auth/login_code_delivery_test.go | 366 ++++++ internal/app/auth/login_email_config_test.go | 3 +- internal/app/auth/login_email_test.go | 28 +- internal/app/auth/premium_grant_test.go | 2 + internal/app/auth/service.go | 606 +++++++-- internal/app/auth/service_test.go | 104 +- internal/app/auth/signup_state_test.go | 554 ++++++++ internal/app/channels/service.go | 78 ++ internal/app/channels/service_test.go | 95 +- internal/app/files/photos.go | 67 +- internal/app/files/seed_test.go | 27 + internal/app/files/service.go | 17 +- internal/app/files/upload_parts_test.go | 45 + internal/app/files/upload_receipt.go | 106 ++ internal/app/maintenance/retention.go | 277 +++- internal/app/maintenance/retention_test.go | 220 +++- internal/app/messages/album_group.go | 28 + internal/app/messages/service.go | 43 + internal/app/messages/service_test.go | 32 + internal/app/updates/service.go | 185 ++- internal/app/updates/service_test.go | 190 ++- internal/compat/layerwire/README.md | 3 +- internal/compat/layerwire/fallback.go | 28 +- internal/compat/layerwire/inbound.go | 59 +- internal/compat/layerwire/routable.go | 80 ++ internal/compat/layerwire/routable_test.go | 43 + .../layerwire/schema/routable-compat.tl | 11 + internal/compat/layerwire/tables.go | 62 +- internal/compat/layerwire/walk.go | 397 +++++- internal/compat/layerwire/walk_limits_test.go | 263 ++++ internal/config/config.go | 104 +- internal/config/config_test.go | 73 +- internal/domain/album_group.go | 51 + internal/domain/channel.go | 75 +- internal/domain/login_code_delivery.go | 55 + internal/domain/media.go | 22 + internal/domain/message.go | 41 +- internal/domain/message_errors.go | 31 +- internal/mtprotoedge/admission.go | 166 +++ internal/mtprotoedge/admission_test.go | 313 +++++ internal/mtprotoedge/auth_key_switch_test.go | 56 + internal/mtprotoedge/conn.go | 50 +- internal/mtprotoedge/encrypted.go | 428 +++++- internal/mtprotoedge/exchange.go | 64 +- internal/mtprotoedge/exchange_compat.go | 75 +- internal/mtprotoedge/exchange_test.go | 285 ++++ internal/mtprotoedge/frame_budget.go | 251 ++++ internal/mtprotoedge/frame_budget_test.go | 337 +++++ internal/mtprotoedge/inbound_rpc.go | 813 ++++++++++-- internal/mtprotoedge/inbound_rpc_test.go | 422 +++++- internal/mtprotoedge/login_email_e2e_test.go | 9 +- internal/mtprotoedge/outbound.go | 1000 ++++++++++++-- internal/mtprotoedge/outbound_scratch.go | 126 ++ internal/mtprotoedge/outbound_test.go | 805 +++++++++++- internal/mtprotoedge/passkey_e2e_test.go | 11 +- .../mtprotoedge/quick_ack_deadline_test.go | 70 + internal/mtprotoedge/rpc_test.go | 146 +++ internal/mtprotoedge/same_port_mux.go | 126 +- internal/mtprotoedge/same_port_mux_test.go | 240 ++++ internal/mtprotoedge/server.go | 315 +++-- internal/mtprotoedge/server_test.go | 2 +- internal/mtprotoedge/session_manager.go | 556 +++++++- internal/mtprotoedge/session_manager_test.go | 519 +++++++- .../session_membership_gen_test.go | 60 +- internal/mtprotoedge/shutdown_gate_test.go | 65 + .../mtprotoedge/structural_limits_test.go | 206 +++ internal/mtprotoedge/transport_compat.go | 257 +++- internal/mtprotoedge/transport_compat_test.go | 14 +- internal/rpc/account.go | 4 +- internal/rpc/account_business.go | 2 +- internal/rpc/account_phone.go | 2 + internal/rpc/account_themes_compat_test.go | 94 +- internal/rpc/album_group.go | 39 + internal/rpc/auth.go | 84 +- internal/rpc/auth_code_rate_limit_test.go | 248 ++++ internal/rpc/auth_qr_login_test.go | 11 +- internal/rpc/bots_longtail.go | 3 +- internal/rpc/channel_fanout_dispatcher.go | 788 +++++++++++- .../rpc/channel_fanout_dispatcher_test.go | 1145 +++++++++++++++++ internal/rpc/channels_legacy_chat.go | 4 +- internal/rpc/channels_members.go | 6 +- internal/rpc/channels_stubs.go | 6 +- internal/rpc/channels_topics.go | 2 +- internal/rpc/channels_updates.go | 8 + internal/rpc/chatlists.go | 4 +- internal/rpc/contacts.go | 18 +- internal/rpc/context.go | 23 + internal/rpc/deps.go | 74 +- internal/rpc/dialogs_rpc_test.go | 9 +- internal/rpc/encrypted_chats.go | 2 +- internal/rpc/errors.go | 6 + internal/rpc/folders.go | 2 +- internal/rpc/idle_backoff.go | 5 +- internal/rpc/idle_backoff_test.go | 31 + internal/rpc/message_idempotency.go | 84 ++ internal/rpc/message_idempotency_test.go | 315 +++++ internal/rpc/message_replay_preflight_test.go | 454 +++++++ internal/rpc/messages_bot_longtail.go | 39 +- .../rpc/messages_bot_longtail_rpc_test.go | 25 +- internal/rpc/messages_compat.go | 2 +- internal/rpc/messages_delete.go | 4 +- internal/rpc/messages_dialogs.go | 18 +- internal/rpc/messages_edit.go | 3 +- internal/rpc/messages_edit_geolive.go | 3 +- internal/rpc/messages_forum_read.go | 2 +- internal/rpc/messages_forward.go | 216 +++- internal/rpc/messages_forward_rpc_test.go | 82 ++ internal/rpc/messages_monoforum.go | 31 +- internal/rpc/messages_pin.go | 8 +- internal/rpc/messages_quick_replies.go | 5 +- internal/rpc/messages_read.go | 2 +- internal/rpc/messages_register.go | 2 +- internal/rpc/messages_saved_dialogs.go | 4 +- internal/rpc/messages_send.go | 140 +- internal/rpc/messages_send_rpc_test.go | 3 + internal/rpc/messages_todos.go | 3 +- internal/rpc/outbox_dispatcher.go | 133 +- internal/rpc/outbox_dispatcher_test.go | 199 ++- internal/rpc/payments_star_gifts.go | 3 +- internal/rpc/phone_conference.go | 17 +- internal/rpc/phone_conference_rpc_test.go | 76 ++ internal/rpc/photos.go | 3 +- internal/rpc/push.go | 6 +- internal/rpc/rate_limit.go | 62 + internal/rpc/request_preflight.go | 137 ++ internal/rpc/request_preflight_test.go | 134 ++ internal/rpc/router.go | 39 +- internal/rpc/router_auth_cache_test.go | 31 +- internal/rpc/rpc_testkit_updates_test.go | 100 +- internal/rpc/send_media.go | 312 +++-- internal/rpc/send_media_test.go | 120 ++ internal/rpc/send_replay.go | 118 ++ internal/rpc/stories.go | 8 +- internal/rpc/temp_key_cache.go | 64 +- internal/rpc/updates.go | 32 +- internal/rpc/updates_rpc_test.go | 121 +- internal/store/album_group.go | 18 + internal/store/bootstrap_update_job.go | 2 + internal/store/channel.go | 11 + internal/store/code.go | 92 +- internal/store/dispatch_outbox.go | 13 +- internal/store/login_code_delivery.go | 67 + internal/store/login_code_delivery_test.go | 67 + internal/store/media.go | 5 + internal/store/memory/album_group.go | 68 + internal/store/memory/album_group_test.go | 131 ++ internal/store/memory/auth.go | 18 +- internal/store/memory/bootstrap_update_job.go | 7 +- internal/store/memory/channel_helpers.go | 10 +- .../store/memory/channel_message_delete.go | 11 +- internal/store/memory/channel_message_edit.go | 10 +- internal/store/memory/channel_message_send.go | 136 +- internal/store/memory/channel_monoforum.go | 49 +- .../memory/channel_monoforum_send_test.go | 21 + .../store/memory/channel_recovery_test.go | 22 + .../memory/channel_send_idempotency_test.go | 76 ++ internal/store/memory/channel_store.go | 117 +- internal/store/memory/channel_test.go | 8 +- .../memory/channel_update_retention_test.go | 150 +++ internal/store/memory/channel_updates.go | 171 ++- internal/store/memory/code_cas.go | 70 + internal/store/memory/code_cas_test.go | 213 +++ internal/store/memory/code_test.go | 24 +- internal/store/memory/login_code.go | 211 +++ internal/store/memory/login_code_delivery.go | 129 ++ .../store/memory/login_code_delivery_test.go | 165 +++ .../memory/login_code_invalidate_test.go | 106 ++ .../store/memory/login_code_state_test.go | 447 +++++++ internal/store/memory/message_delete.go | 22 +- internal/store/memory/message_helpers.go | 18 +- .../store/memory/message_idempotency_test.go | 164 +++ internal/store/memory/message_send.go | 137 +- internal/store/memory/message_store.go | 36 +- internal/store/memory/message_test.go | 10 +- .../store/memory/scoped_code_state_test.go | 240 ++++ internal/store/memory/updates.go | 41 +- internal/store/postgres/album_group.go | 162 +++ .../postgres/album_group_integration_test.go | 118 ++ .../postgres/auth_login_integration_test.go | 25 +- internal/store/postgres/authkey.go | 152 ++- .../authkey_retention_integration_test.go | 215 ++++ internal/store/postgres/authorization.go | 132 +- .../store/postgres/bootstrap_update_job.go | 2 +- .../bootstrap_update_job_integration_test.go | 53 + .../channel_difference_integration_test.go | 64 +- internal/store/postgres/channel_helpers.go | 9 +- .../store/postgres/channel_message_delete.go | 10 +- .../store/postgres/channel_message_send.go | 173 ++- internal/store/postgres/channel_monoforum.go | 72 +- ...channel_monoforum_send_integration_test.go | 67 + ...annel_send_idempotency_integration_test.go | 490 +++++++ .../postgres/channel_update_retention.go | 285 ++++ ...annel_update_retention_integration_test.go | 232 ++++ internal/store/postgres/channel_updates.go | 42 +- .../contiguous_pts_integration_test.go | 7 +- internal/store/postgres/dispatch_outbox.go | 127 +- ...spatch_outbox_sharding_integration_test.go | 352 +++++ .../postgres/groupcall_integration_test.go | 20 +- .../store/postgres/login_code_delivery.go | 342 +++++ .../login_code_delivery_integration_test.go | 475 +++++++ .../postgres/login_code_delivery_retention.go | 32 + internal/store/postgres/media.go | 63 + .../store/postgres/media_integration_test.go | 37 + internal/store/postgres/message_delete.go | 23 + .../message_delete_integration_test.go | 14 +- .../message_markup_integration_test.go | 22 +- .../postgres/message_plan_integration_test.go | 46 +- internal/store/postgres/message_send.go | 234 +++- ...ssage_send_idempotency_integration_test.go | 319 +++++ .../postgres/message_send_idempotency_test.go | 83 ++ .../postgres/message_send_integration_test.go | 11 +- ..._send_migration_compat_integration_test.go | 222 ++++ internal/store/postgres/queries/message.sql | 33 +- .../postgres/queries/user_update_event.sql | 158 ++- .../retention_migration_integration_test.go | 137 ++ .../postgres/saved_dialog_integration_test.go | 14 +- .../send_replay_migration_integration_test.go | 45 + .../store/postgres/sqlcgen/message.sql.go | 171 ++- internal/store/postgres/sqlcgen/models.go | 130 +- .../postgres/sqlcgen/user_update_event.sql.go | 252 +++- .../store/postgres/update_event_retention.go | 300 +++++ ...update_event_retention_integration_test.go | 643 +++++++++ internal/store/postgres/updatestate.go | 20 +- internal/store/private_send_idempotency.go | 243 ++++ .../store/private_send_idempotency_test.go | 56 + .../redisstore/allocator_integration_test.go | 17 + internal/store/redisstore/code.go | 51 +- internal/store/redisstore/code_cas.go | 147 +++ .../redisstore/code_cas_integration_test.go | 237 ++++ .../store/redisstore/code_integration_test.go | 32 + internal/store/redisstore/login_code.go | 385 ++++++ .../redisstore/login_code_integration_test.go | 517 ++++++++ .../login_code_invalidate_integration_test.go | 103 ++ internal/store/redisstore/ratelimit.go | 44 +- .../scoped_code_state_integration_test.go | 288 +++++ internal/store/send_replay.go | 20 + internal/store/storetest/groupcall.go | 46 +- internal/store/updatestate.go | 4 + 283 files changed, 29231 insertions(+), 2295 deletions(-) create mode 100644 deploy/migrations/0062_private_message_idempotency.down.sql create mode 100644 deploy/migrations/0062_private_message_idempotency.up.sql create mode 100644 deploy/migrations/0063_channel_update_retention.down.sql create mode 100644 deploy/migrations/0063_channel_update_retention.up.sql create mode 100644 deploy/migrations/0064_user_update_retention.down.sql create mode 100644 deploy/migrations/0064_user_update_retention.up.sql create mode 100644 deploy/migrations/0065_auth_key_orphan_retention.down.sql create mode 100644 deploy/migrations/0065_auth_key_orphan_retention.up.sql create mode 100644 deploy/migrations/0066_dispatch_outbox_user_lanes.down.sql create mode 100644 deploy/migrations/0066_dispatch_outbox_user_lanes.up.sql create mode 100644 deploy/migrations/0067_auth_key_last_used.down.sql create mode 100644 deploy/migrations/0067_auth_key_last_used.up.sql create mode 100644 deploy/migrations/0068_private_message_send_receipt.down.sql create mode 100644 deploy/migrations/0068_private_message_send_receipt.up.sql create mode 100644 deploy/migrations/0069_dispatch_outbox_user_heads.down.sql create mode 100644 deploy/migrations/0069_dispatch_outbox_user_heads.up.sql create mode 100644 deploy/migrations/0070_dispatch_outbox_head_readiness.down.sql create mode 100644 deploy/migrations/0070_dispatch_outbox_head_readiness.up.sql create mode 100644 deploy/migrations/0071_temp_auth_key_expiry_seek.down.sql create mode 100644 deploy/migrations/0071_temp_auth_key_expiry_seek.up.sql create mode 100644 deploy/migrations/0072_dispatch_outbox_head_ready_indexes.down.sql create mode 100644 deploy/migrations/0072_dispatch_outbox_head_ready_indexes.up.sql create mode 100644 deploy/migrations/0073_send_replay_snapshots.down.sql create mode 100644 deploy/migrations/0073_send_replay_snapshots.up.sql create mode 100644 deploy/migrations/0074_dispatch_outbox_poison_head_index.down.sql create mode 100644 deploy/migrations/0074_dispatch_outbox_poison_head_index.up.sql create mode 100644 deploy/migrations/0075_uploaded_media_receipts.down.sql create mode 100644 deploy/migrations/0075_uploaded_media_receipts.up.sql create mode 100644 deploy/migrations/0076_album_group_reservations.down.sql create mode 100644 deploy/migrations/0076_album_group_reservations.up.sql create mode 100644 deploy/migrations/0077_correct_private_send_defaults.down.sql create mode 100644 deploy/migrations/0077_correct_private_send_defaults.up.sql create mode 100644 deploy/migrations/0078_channel_send_fingerprint.down.sql create mode 100644 deploy/migrations/0078_channel_send_fingerprint.up.sql create mode 100644 deploy/migrations/0079_login_code_message_deliveries.down.sql create mode 100644 deploy/migrations/0079_login_code_message_deliveries.up.sql create mode 100644 deploy/migrations/0080_login_code_delivery_expiry.down.sql create mode 100644 deploy/migrations/0080_login_code_delivery_expiry.up.sql create mode 100644 internal/app/auth/login_code_delivery_test.go create mode 100644 internal/app/auth/signup_state_test.go create mode 100644 internal/app/files/upload_receipt.go create mode 100644 internal/app/messages/album_group.go create mode 100644 internal/compat/layerwire/routable.go create mode 100644 internal/compat/layerwire/routable_test.go create mode 100644 internal/compat/layerwire/schema/routable-compat.tl create mode 100644 internal/compat/layerwire/walk_limits_test.go create mode 100644 internal/domain/album_group.go create mode 100644 internal/domain/login_code_delivery.go create mode 100644 internal/mtprotoedge/admission.go create mode 100644 internal/mtprotoedge/admission_test.go create mode 100644 internal/mtprotoedge/auth_key_switch_test.go create mode 100644 internal/mtprotoedge/frame_budget.go create mode 100644 internal/mtprotoedge/frame_budget_test.go create mode 100644 internal/mtprotoedge/outbound_scratch.go create mode 100644 internal/mtprotoedge/quick_ack_deadline_test.go create mode 100644 internal/mtprotoedge/same_port_mux_test.go create mode 100644 internal/mtprotoedge/shutdown_gate_test.go create mode 100644 internal/mtprotoedge/structural_limits_test.go create mode 100644 internal/rpc/album_group.go create mode 100644 internal/rpc/auth_code_rate_limit_test.go create mode 100644 internal/rpc/message_idempotency.go create mode 100644 internal/rpc/message_idempotency_test.go create mode 100644 internal/rpc/message_replay_preflight_test.go create mode 100644 internal/rpc/request_preflight.go create mode 100644 internal/rpc/request_preflight_test.go create mode 100644 internal/rpc/send_replay.go create mode 100644 internal/store/album_group.go create mode 100644 internal/store/login_code_delivery.go create mode 100644 internal/store/login_code_delivery_test.go create mode 100644 internal/store/memory/album_group.go create mode 100644 internal/store/memory/album_group_test.go create mode 100644 internal/store/memory/channel_recovery_test.go create mode 100644 internal/store/memory/channel_send_idempotency_test.go create mode 100644 internal/store/memory/channel_update_retention_test.go create mode 100644 internal/store/memory/code_cas.go create mode 100644 internal/store/memory/code_cas_test.go create mode 100644 internal/store/memory/login_code.go create mode 100644 internal/store/memory/login_code_delivery.go create mode 100644 internal/store/memory/login_code_delivery_test.go create mode 100644 internal/store/memory/login_code_invalidate_test.go create mode 100644 internal/store/memory/login_code_state_test.go create mode 100644 internal/store/memory/message_idempotency_test.go create mode 100644 internal/store/memory/scoped_code_state_test.go create mode 100644 internal/store/postgres/album_group.go create mode 100644 internal/store/postgres/album_group_integration_test.go create mode 100644 internal/store/postgres/authkey_retention_integration_test.go create mode 100644 internal/store/postgres/bootstrap_update_job_integration_test.go create mode 100644 internal/store/postgres/channel_send_idempotency_integration_test.go create mode 100644 internal/store/postgres/channel_update_retention.go create mode 100644 internal/store/postgres/channel_update_retention_integration_test.go create mode 100644 internal/store/postgres/dispatch_outbox_sharding_integration_test.go create mode 100644 internal/store/postgres/login_code_delivery.go create mode 100644 internal/store/postgres/login_code_delivery_integration_test.go create mode 100644 internal/store/postgres/login_code_delivery_retention.go create mode 100644 internal/store/postgres/message_send_idempotency_integration_test.go create mode 100644 internal/store/postgres/message_send_idempotency_test.go create mode 100644 internal/store/postgres/private_send_migration_compat_integration_test.go create mode 100644 internal/store/postgres/retention_migration_integration_test.go create mode 100644 internal/store/postgres/send_replay_migration_integration_test.go create mode 100644 internal/store/postgres/update_event_retention.go create mode 100644 internal/store/postgres/update_event_retention_integration_test.go create mode 100644 internal/store/private_send_idempotency.go create mode 100644 internal/store/private_send_idempotency_test.go create mode 100644 internal/store/redisstore/code_cas.go create mode 100644 internal/store/redisstore/code_cas_integration_test.go create mode 100644 internal/store/redisstore/login_code.go create mode 100644 internal/store/redisstore/login_code_integration_test.go create mode 100644 internal/store/redisstore/login_code_invalidate_integration_test.go create mode 100644 internal/store/redisstore/scoped_code_state_integration_test.go create mode 100644 internal/store/send_replay.go diff --git a/.env.example b/.env.example index b5f7b2c7..30703da4 100644 --- a/.env.example +++ b/.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 diff --git a/cmd/telesrv/main.go b/cmd/telesrv/main.go index d098e9dd..e7c151cc 100644 --- a/cmd/telesrv/main.go +++ b/cmd/telesrv/main.go @@ -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), diff --git a/deploy/migrations/0062_private_message_idempotency.down.sql b/deploy/migrations/0062_private_message_idempotency.down.sql new file mode 100644 index 00000000..ea3efaaa --- /dev/null +++ b/deploy/migrations/0062_private_message_idempotency.down.sql @@ -0,0 +1,3 @@ +ALTER TABLE private_messages + DROP COLUMN IF EXISTS recipient_delivered, + DROP COLUMN IF EXISTS request_fingerprint; diff --git a/deploy/migrations/0062_private_message_idempotency.up.sql b/deploy/migrations/0062_private_message_idempotency.up.sql new file mode 100644 index 00000000..29487a3b --- /dev/null +++ b/deploy/migrations/0062_private_message_idempotency.up.sql @@ -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. diff --git a/deploy/migrations/0063_channel_update_retention.down.sql b/deploy/migrations/0063_channel_update_retention.down.sql new file mode 100644 index 00000000..36e82d78 --- /dev/null +++ b/deploy/migrations/0063_channel_update_retention.down.sql @@ -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; diff --git a/deploy/migrations/0063_channel_update_retention.up.sql b/deploy/migrations/0063_channel_update_retention.up.sql new file mode 100644 index 00000000..bd88cb2e --- /dev/null +++ b/deploy/migrations/0063_channel_update_retention.up.sql @@ -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); diff --git a/deploy/migrations/0064_user_update_retention.down.sql b/deploy/migrations/0064_user_update_retention.down.sql new file mode 100644 index 00000000..dd9532f7 --- /dev/null +++ b/deploy/migrations/0064_user_update_retention.down.sql @@ -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; diff --git a/deploy/migrations/0064_user_update_retention.up.sql b/deploy/migrations/0064_user_update_retention.up.sql new file mode 100644 index 00000000..ca882b5e --- /dev/null +++ b/deploy/migrations/0064_user_update_retention.up.sql @@ -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); diff --git a/deploy/migrations/0065_auth_key_orphan_retention.down.sql b/deploy/migrations/0065_auth_key_orphan_retention.down.sql new file mode 100644 index 00000000..fd2c6ce4 --- /dev/null +++ b/deploy/migrations/0065_auth_key_orphan_retention.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS public.auth_keys_orphan_retention_idx; diff --git a/deploy/migrations/0065_auth_key_orphan_retention.up.sql b/deploy/migrations/0065_auth_key_orphan_retention.up.sql new file mode 100644 index 00000000..a7c38043 --- /dev/null +++ b/deploy/migrations/0065_auth_key_orphan_retention.up.sql @@ -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); diff --git a/deploy/migrations/0066_dispatch_outbox_user_lanes.down.sql b/deploy/migrations/0066_dispatch_outbox_user_lanes.down.sql new file mode 100644 index 00000000..265e48f1 --- /dev/null +++ b/deploy/migrations/0066_dispatch_outbox_user_lanes.down.sql @@ -0,0 +1,2 @@ +DROP INDEX IF EXISTS dispatch_outbox_logical_shard_head_idx; +DROP INDEX IF EXISTS dispatch_outbox_user_pts_uidx; diff --git a/deploy/migrations/0066_dispatch_outbox_user_lanes.up.sql b/deploy/migrations/0066_dispatch_outbox_user_lanes.up.sql new file mode 100644 index 00000000..a669d907 --- /dev/null +++ b/deploy/migrations/0066_dispatch_outbox_user_lanes.up.sql @@ -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 + ); diff --git a/deploy/migrations/0067_auth_key_last_used.down.sql b/deploy/migrations/0067_auth_key_last_used.down.sql new file mode 100644 index 00000000..dd5f0053 --- /dev/null +++ b/deploy/migrations/0067_auth_key_last_used.down.sql @@ -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); diff --git a/deploy/migrations/0067_auth_key_last_used.up.sql b/deploy/migrations/0067_auth_key_last_used.up.sql new file mode 100644 index 00000000..2ef10380 --- /dev/null +++ b/deploy/migrations/0067_auth_key_last_used.up.sql @@ -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); diff --git a/deploy/migrations/0068_private_message_send_receipt.down.sql b/deploy/migrations/0068_private_message_send_receipt.down.sql new file mode 100644 index 00000000..f98d2b2a --- /dev/null +++ b/deploy/migrations/0068_private_message_send_receipt.down.sql @@ -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; diff --git a/deploy/migrations/0068_private_message_send_receipt.up.sql b/deploy/migrations/0068_private_message_send_receipt.up.sql new file mode 100644 index 00000000..6b249d6f --- /dev/null +++ b/deploy/migrations/0068_private_message_send_receipt.up.sql @@ -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. diff --git a/deploy/migrations/0069_dispatch_outbox_user_heads.down.sql b/deploy/migrations/0069_dispatch_outbox_user_heads.down.sql new file mode 100644 index 00000000..1a577420 --- /dev/null +++ b/deploy/migrations/0069_dispatch_outbox_user_heads.down.sql @@ -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; diff --git a/deploy/migrations/0069_dispatch_outbox_user_heads.up.sql b/deploy/migrations/0069_dispatch_outbox_user_heads.up.sql new file mode 100644 index 00000000..39c0c993 --- /dev/null +++ b/deploy/migrations/0069_dispatch_outbox_user_heads.up.sql @@ -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; diff --git a/deploy/migrations/0070_dispatch_outbox_head_readiness.down.sql b/deploy/migrations/0070_dispatch_outbox_head_readiness.down.sql new file mode 100644 index 00000000..59f46793 --- /dev/null +++ b/deploy/migrations/0070_dispatch_outbox_head_readiness.down.sql @@ -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; diff --git a/deploy/migrations/0070_dispatch_outbox_head_readiness.up.sql b/deploy/migrations/0070_dispatch_outbox_head_readiness.up.sql new file mode 100644 index 00000000..82f4d3c0 --- /dev/null +++ b/deploy/migrations/0070_dispatch_outbox_head_readiness.up.sql @@ -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(); diff --git a/deploy/migrations/0071_temp_auth_key_expiry_seek.down.sql b/deploy/migrations/0071_temp_auth_key_expiry_seek.down.sql new file mode 100644 index 00000000..e0692ea3 --- /dev/null +++ b/deploy/migrations/0071_temp_auth_key_expiry_seek.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS public.temp_auth_key_bindings_expiry_idx; diff --git a/deploy/migrations/0071_temp_auth_key_expiry_seek.up.sql b/deploy/migrations/0071_temp_auth_key_expiry_seek.up.sql new file mode 100644 index 00000000..50d48939 --- /dev/null +++ b/deploy/migrations/0071_temp_auth_key_expiry_seek.up.sql @@ -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); diff --git a/deploy/migrations/0072_dispatch_outbox_head_ready_indexes.down.sql b/deploy/migrations/0072_dispatch_outbox_head_ready_indexes.down.sql new file mode 100644 index 00000000..ed866875 --- /dev/null +++ b/deploy/migrations/0072_dispatch_outbox_head_ready_indexes.down.sql @@ -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; diff --git a/deploy/migrations/0072_dispatch_outbox_head_ready_indexes.up.sql b/deploy/migrations/0072_dispatch_outbox_head_ready_indexes.up.sql new file mode 100644 index 00000000..4af9837e --- /dev/null +++ b/deploy/migrations/0072_dispatch_outbox_head_ready_indexes.up.sql @@ -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; diff --git a/deploy/migrations/0073_send_replay_snapshots.down.sql b/deploy/migrations/0073_send_replay_snapshots.down.sql new file mode 100644 index 00000000..6309617d --- /dev/null +++ b/deploy/migrations/0073_send_replay_snapshots.down.sql @@ -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; diff --git a/deploy/migrations/0073_send_replay_snapshots.up.sql b/deploy/migrations/0073_send_replay_snapshots.up.sql new file mode 100644 index 00000000..84671977 --- /dev/null +++ b/deploy/migrations/0073_send_replay_snapshots.up.sql @@ -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. diff --git a/deploy/migrations/0074_dispatch_outbox_poison_head_index.down.sql b/deploy/migrations/0074_dispatch_outbox_poison_head_index.down.sql new file mode 100644 index 00000000..4a6ab770 --- /dev/null +++ b/deploy/migrations/0074_dispatch_outbox_poison_head_index.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS public.dispatch_outbox_user_heads_failed_cleanup_idx; diff --git a/deploy/migrations/0074_dispatch_outbox_poison_head_index.up.sql b/deploy/migrations/0074_dispatch_outbox_poison_head_index.up.sql new file mode 100644 index 00000000..89a47e59 --- /dev/null +++ b/deploy/migrations/0074_dispatch_outbox_poison_head_index.up.sql @@ -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'; diff --git a/deploy/migrations/0075_uploaded_media_receipts.down.sql b/deploy/migrations/0075_uploaded_media_receipts.down.sql new file mode 100644 index 00000000..ad94ebb5 --- /dev/null +++ b/deploy/migrations/0075_uploaded_media_receipts.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS public.uploaded_media_receipts; diff --git a/deploy/migrations/0075_uploaded_media_receipts.up.sql b/deploy/migrations/0075_uploaded_media_receipts.up.sql new file mode 100644 index 00000000..7f61dfcc --- /dev/null +++ b/deploy/migrations/0075_uploaded_media_receipts.up.sql @@ -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) +); diff --git a/deploy/migrations/0076_album_group_reservations.down.sql b/deploy/migrations/0076_album_group_reservations.down.sql new file mode 100644 index 00000000..f4db22cc --- /dev/null +++ b/deploy/migrations/0076_album_group_reservations.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS album_group_reservations; diff --git a/deploy/migrations/0076_album_group_reservations.up.sql b/deploy/migrations/0076_album_group_reservations.up.sql new file mode 100644 index 00000000..b58c76bc --- /dev/null +++ b/deploy/migrations/0076_album_group_reservations.up.sql @@ -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.'; diff --git a/deploy/migrations/0077_correct_private_send_defaults.down.sql b/deploy/migrations/0077_correct_private_send_defaults.down.sql new file mode 100644 index 00000000..dea274ec --- /dev/null +++ b/deploy/migrations/0077_correct_private_send_defaults.down.sql @@ -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; diff --git a/deploy/migrations/0077_correct_private_send_defaults.up.sql b/deploy/migrations/0077_correct_private_send_defaults.up.sql new file mode 100644 index 00000000..ea557db2 --- /dev/null +++ b/deploy/migrations/0077_correct_private_send_defaults.up.sql @@ -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. diff --git a/deploy/migrations/0078_channel_send_fingerprint.down.sql b/deploy/migrations/0078_channel_send_fingerprint.down.sql new file mode 100644 index 00000000..efc62acc --- /dev/null +++ b/deploy/migrations/0078_channel_send_fingerprint.down.sql @@ -0,0 +1,3 @@ +ALTER TABLE public.channel_messages + DROP CONSTRAINT IF EXISTS channel_messages_request_fingerprint_size, + DROP COLUMN IF EXISTS request_fingerprint; diff --git a/deploy/migrations/0078_channel_send_fingerprint.up.sql b/deploy/migrations/0078_channel_send_fingerprint.up.sql new file mode 100644 index 00000000..56bf1e51 --- /dev/null +++ b/deploy/migrations/0078_channel_send_fingerprint.up.sql @@ -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. diff --git a/deploy/migrations/0079_login_code_message_deliveries.down.sql b/deploy/migrations/0079_login_code_message_deliveries.down.sql new file mode 100644 index 00000000..ebc7e61e --- /dev/null +++ b/deploy/migrations/0079_login_code_message_deliveries.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS login_code_message_deliveries; diff --git a/deploy/migrations/0079_login_code_message_deliveries.up.sql b/deploy/migrations/0079_login_code_message_deliveries.up.sql new file mode 100644 index 00000000..6f853ac3 --- /dev/null +++ b/deploy/migrations/0079_login_code_message_deliveries.up.sql @@ -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'; diff --git a/deploy/migrations/0080_login_code_delivery_expiry.down.sql b/deploy/migrations/0080_login_code_delivery_expiry.down.sql new file mode 100644 index 00000000..e6ff0813 --- /dev/null +++ b/deploy/migrations/0080_login_code_delivery_expiry.down.sql @@ -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; diff --git a/deploy/migrations/0080_login_code_delivery_expiry.up.sql b/deploy/migrations/0080_login_code_delivery_expiry.up.sql new file mode 100644 index 00000000..23cbd8be --- /dev/null +++ b/deploy/migrations/0080_login_code_delivery_expiry.up.sql @@ -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); diff --git a/internal/app/account/login_email_test.go b/internal/app/account/login_email_test.go index b0e825fc..5871a469 100644 --- a/internal/app/account/login_email_test.go +++ b/internal/app/account/login_email_test.go @@ -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) + } +} diff --git a/internal/app/account/phone_change.go b/internal/app/account/phone_change.go index cc559f27..f69015dc 100644 --- a/internal/app/account/phone_change.go +++ b/internal/app/account/phone_change.go @@ -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 { diff --git a/internal/app/account/phone_change_test.go b/internal/app/account/phone_change_test.go index e6987dab..f541e267 100644 --- a/internal/app/account/phone_change_test.go +++ b/internal/app/account/phone_change_test.go @@ -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 }() } diff --git a/internal/app/account/service.go b/internal/app/account/service.go index 9c39fd38..78ba8a2a 100644 --- a/internal/app/account/service.go +++ b/internal/app/account/service.go @@ -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 { diff --git a/internal/app/auth/change_phone_resend_test.go b/internal/app/auth/change_phone_resend_test.go index b139768f..ccf6ffac 100644 --- a/internal/app/auth/change_phone_resend_test.go +++ b/internal/app/auth/change_phone_resend_test.go @@ -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") + } +} diff --git a/internal/app/auth/login_code_delivery_test.go b/internal/app/auth/login_code_delivery_test.go new file mode 100644 index 00000000..21666bfb --- /dev/null +++ b/internal/app/auth/login_code_delivery_test.go @@ -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) + } +} diff --git a/internal/app/auth/login_email_config_test.go b/internal/app/auth/login_email_config_test.go index 3de25ec6..9b9271aa 100644 --- a/internal/app/auth/login_email_config_test.go +++ b/internal/app/auth/login_email_config_test.go @@ -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 } diff --git a/internal/app/auth/login_email_test.go b/internal/app/auth/login_email_test.go index dccd5113..8fe0e969 100644 --- a/internal/app/auth/login_email_test.go +++ b/internal/app/auth/login_email_test.go @@ -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) } diff --git a/internal/app/auth/premium_grant_test.go b/internal/app/auth/premium_grant_test.go index 7ddfb635..665028d0 100644 --- a/internal/app/auth/premium_grant_test.go +++ b/internal/app/auth/premium_grant_test.go @@ -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) diff --git a/internal/app/auth/service.go b/internal/app/auth/service.go index e18756ed..d332fb70 100644 --- a/internal/app/auth/service.go +++ b/internal/app/auth/service.go @@ -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) } diff --git a/internal/app/auth/service_test.go b/internal/app/auth/service_test.go index c37bdff4..73a8b562 100644 --- a/internal/app/auth/service_test.go +++ b/internal/app/auth/service_test.go @@ -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 { diff --git a/internal/app/auth/signup_state_test.go b/internal/app/auth/signup_state_test.go new file mode 100644 index 00000000..1b369885 --- /dev/null +++ b/internal/app/auth/signup_state_test.go @@ -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) + } +} diff --git a/internal/app/channels/service.go b/internal/app/channels/service.go index 7f075b7b..6ce6f327 100644 --- a/internal/app/channels/service.go +++ b/internal/app/channels/service.go @@ -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 { diff --git a/internal/app/channels/service_test.go b/internal/app/channels/service_test.go index e463eca2..34333bb5 100644 --- a/internal/app/channels/service_test.go +++ b/internal/app/channels/service_test.go @@ -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{ diff --git a/internal/app/files/photos.go b/internal/app/files/photos.go index 6a215fe6..e5276e8c 100644 --- a/internal/app/files/photos.go +++ b/internal/app/files/photos.go @@ -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))] { diff --git a/internal/app/files/seed_test.go b/internal/app/files/seed_test.go index 023bf6b3..2b53614b 100644 --- a/internal/app/files/seed_test.go +++ b/internal/app/files/seed_test.go @@ -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() diff --git a/internal/app/files/service.go b/internal/app/files/service.go index fc3b74b5..135a4f32 100644 --- a/internal/app/files/service.go +++ b/internal/app/files/service.go @@ -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 } diff --git a/internal/app/files/upload_parts_test.go b/internal/app/files/upload_parts_test.go index fc764739..fac32ebf 100644 --- a/internal/app/files/upload_parts_test.go +++ b/internal/app/files/upload_parts_test.go @@ -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 { diff --git a/internal/app/files/upload_receipt.go b/internal/app/files/upload_receipt.go new file mode 100644 index 00000000..f046d3b2 --- /dev/null +++ b/internal/app/files/upload_receipt.go @@ -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), + ) + } +} diff --git a/internal/app/maintenance/retention.go b/internal/app/maintenance/retention.go index fdbb54a8..78736c10 100644 --- a/internal/app/maintenance/retention.go +++ b/internal/app/maintenance/retention.go @@ -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 } diff --git a/internal/app/maintenance/retention_test.go b/internal/app/maintenance/retention_test.go index 14da5fdd..35e72f85 100644 --- a/internal/app/maintenance/retention_test.go +++ b/internal/app/maintenance/retention_test.go @@ -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) + } +} diff --git a/internal/app/messages/album_group.go b/internal/app/messages/album_group.go new file mode 100644 index 00000000..7c653ccb --- /dev/null +++ b/internal/app/messages/album_group.go @@ -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) +} diff --git a/internal/app/messages/service.go b/internal/app/messages/service.go index 0d2f305d..f5bd75d9 100644 --- a/internal/app/messages/service.go +++ b/internal/app/messages/service.go @@ -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 diff --git a/internal/app/messages/service_test.go b/internal/app/messages/service_test.go index 9511bde0..ef33549c 100644 --- a/internal/app/messages/service_test.go +++ b/internal/app/messages/service_test.go @@ -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{} diff --git a/internal/app/updates/service.go b/internal/app/updates/service.go index 91f11722..889488ee 100644 --- a/internal/app/updates/service.go +++ b/internal/app/updates/service.go @@ -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 } } diff --git a/internal/app/updates/service_test.go b/internal/app/updates/service_test.go index ad118dea..306ac495 100644 --- a/internal/app/updates/service_test.go +++ b/internal/app/updates/service_test.go @@ -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)) diff --git a/internal/compat/layerwire/README.md b/internal/compat/layerwire/README.md index 0d298177..4dd10d49 100644 --- a/internal/compat/layerwire/README.md +++ b/internal/compat/layerwire/README.md @@ -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->Vector`）。 - 报 `field X not defaultable` 或字段**改名** → 往 `inbound.go driftFieldRenames` 加 `"\x00<227字段>": "<老字段>"`（参照 `bots.exportBotToken\x00bot`）。 -4. **绝不**为此写一个新的 `handleLegacyXxx` 解码 handler——那是旧做法，已全删。统一走数据 + 通用引擎。 +4. **绝不**为此写一个新的 `handleLegacyXxx` 解码 handler——统一走数据 + 通用引擎。`routable-compat.tl` 只给既存 DrKLO theme fallback 补 dispatcher 前结构门禁，不是新增 adapter 的入口。 ## 操作 4：出站 `TestCoverageGate` 失败 diff --git a/internal/compat/layerwire/fallback.go b/internal/compat/layerwire/fallback.go index 7c7fc4aa..b9bb043e 100644 --- a/internal/compat/layerwire/fallback.go +++ b/internal/compat/layerwire/fallback.go @@ -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); // 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 // <=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 } } diff --git a/internal/compat/layerwire/inbound.go b/internal/compat/layerwire/inbound.go index de413c91..2cc7bd5f 100644 --- a/internal/compat/layerwire/inbound.go +++ b/internal/compat/layerwire/inbound.go @@ -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 "->"; 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 -> id:Vector (wrap each int in inputMessageID). - "Vector->Vector": func(raw []byte, out *bin.Buffer) error { + "Vector->Vector": 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 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 { diff --git a/internal/compat/layerwire/routable.go b/internal/compat/layerwire/routable.go new file mode 100644 index 00000000..3f17984e --- /dev/null +++ b/internal/compat/layerwire/routable.go @@ -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) +} diff --git a/internal/compat/layerwire/routable_test.go b/internal/compat/layerwire/routable_test.go new file mode 100644 index 00000000..5c4e8010 --- /dev/null +++ b/internal/compat/layerwire/routable_test.go @@ -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) + } + }) +} diff --git a/internal/compat/layerwire/schema/routable-compat.tl b/internal/compat/layerwire/schema/routable-compat.tl new file mode 100644 index 00000000..dd5aaee4 --- /dev/null +++ b/internal/compat/layerwire/schema/routable-compat.tl @@ -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; diff --git a/internal/compat/layerwire/tables.go b/internal/compat/layerwire/tables.go index a0b59079..4290d04a 100644 --- a/internal/compat/layerwire/tables.go +++ b/internal/compat/layerwire/tables.go @@ -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) diff --git a/internal/compat/layerwire/walk.go b/internal/compat/layerwire/walk.go index 3d9d7df9..9f89aff3 100644 --- a/internal/compat/layerwire/walk.go +++ b/internal/compat/layerwire/walk.go @@ -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< 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 "" + } + return cl.name +} + +func fieldName(f *fieldLayout) string { + if f == nil || f.name == "" { + return "" + } + return f.name +} diff --git a/internal/compat/layerwire/walk_limits_test.go b/internal/compat/layerwire/walk_limits_test.go new file mode 100644 index 00000000..ca26f885 --- /dev/null +++ b/internal/compat/layerwire/walk_limits_test.go @@ -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) + } + }) +} diff --git a/internal/config/config.go b/internal/config/config.go index fdd7ede8..0f77365d 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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), diff --git a/internal/config/config_test.go b/internal/config/config_test.go index 288c4a89..783058a2 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -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) diff --git a/internal/domain/album_group.go b/internal/domain/album_group.go new file mode 100644 index 00000000..9ab6fa63 --- /dev/null +++ b/internal/domain/album_group.go @@ -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 +} diff --git a/internal/domain/channel.go b/internal/domain/channel.go index 418fe7b9..504b7239 100644 --- a/internal/domain/channel.go +++ b/internal/domain/channel.go @@ -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 diff --git a/internal/domain/login_code_delivery.go b/internal/domain/login_code_delivery.go new file mode 100644 index 00000000..d8f3a428 --- /dev/null +++ b/internal/domain/login_code_delivery.go @@ -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 +} diff --git a/internal/domain/media.go b/internal/domain/media.go index a19a1d58..8c898f1c 100644 --- a/internal/domain/media.go +++ b/internal/domain/media.go @@ -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 diff --git a/internal/domain/message.go b/internal/domain/message.go index 5b23f5a9..443a224a 100644 --- a/internal/domain/message.go +++ b/internal/domain/message.go @@ -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 命令。 diff --git a/internal/domain/message_errors.go b/internal/domain/message_errors.go index 8a1e49f7..342c0ba2 100644 --- a/internal/domain/message_errors.go +++ b/internal/domain/message_errors.go @@ -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") diff --git a/internal/mtprotoedge/admission.go b/internal/mtprotoedge/admission.go new file mode 100644 index 00000000..8b2506a3 --- /dev/null +++ b/internal/mtprotoedge/admission.go @@ -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 "" + } + 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 + } +} diff --git a/internal/mtprotoedge/admission_test.go b/internal/mtprotoedge/admission_test.go new file mode 100644 index 00000000..1ce201ef --- /dev/null +++ b/internal/mtprotoedge/admission_test.go @@ -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) + } +} diff --git a/internal/mtprotoedge/auth_key_switch_test.go b/internal/mtprotoedge/auth_key_switch_test.go new file mode 100644 index 00000000..5bd1e65d --- /dev/null +++ b/internal/mtprotoedge/auth_key_switch_test.go @@ -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) + } +} diff --git a/internal/mtprotoedge/conn.go b/internal/mtprotoedge/conn.go index 53e4ce6d..2634711f 100644 --- a/internal/mtprotoedge/conn.go +++ b/internal/mtprotoedge/conn.go @@ -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 diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index e514fc49..fceb4244 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -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) diff --git a/internal/mtprotoedge/exchange.go b/internal/mtprotoedge/exchange.go index f8ce0366..04221f87 100644 --- a/internal/mtprotoedge/exchange.go +++ b/internal/mtprotoedge/exchange.go @@ -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()} -} diff --git a/internal/mtprotoedge/exchange_compat.go b/internal/mtprotoedge/exchange_compat.go index 0430f5f8..24b16730 100644 --- a/internal/mtprotoedge/exchange_compat.go +++ b/internal/mtprotoedge/exchange_compat.go @@ -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, } } diff --git a/internal/mtprotoedge/exchange_test.go b/internal/mtprotoedge/exchange_test.go index 2d875262..3d9825a8 100644 --- a/internal/mtprotoedge/exchange_test.go +++ b/internal/mtprotoedge/exchange_test.go @@ -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 diff --git a/internal/mtprotoedge/frame_budget.go b/internal/mtprotoedge/frame_budget.go new file mode 100644 index 00000000..69e11379 --- /dev/null +++ b/internal/mtprotoedge/frame_budget.go @@ -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) +} diff --git a/internal/mtprotoedge/frame_budget_test.go b/internal/mtprotoedge/frame_budget_test.go new file mode 100644 index 00000000..84182f8a --- /dev/null +++ b/internal/mtprotoedge/frame_budget_test.go @@ -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() +} diff --git a/internal/mtprotoedge/inbound_rpc.go b/internal/mtprotoedge/inbound_rpc.go index 15d05ece..c6bafed8 100644 --- a/internal/mtprotoedge/inbound_rpc.go +++ b/internal/mtprotoedge/inbound_rpc.go @@ -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 } } diff --git a/internal/mtprotoedge/inbound_rpc_test.go b/internal/mtprotoedge/inbound_rpc_test.go index c085b9ef..d26a2922 100644 --- a/internal/mtprotoedge/inbound_rpc_test.go +++ b/internal/mtprotoedge/inbound_rpc_test.go @@ -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) + } } diff --git a/internal/mtprotoedge/login_email_e2e_test.go b/internal/mtprotoedge/login_email_e2e_test.go index cd486d45..e59003cc 100644 --- a/internal/mtprotoedge/login_email_e2e_test.go +++ b/internal/mtprotoedge/login_email_e2e_test.go @@ -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}) diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index 53359c24..cc737833 100644 --- a/internal/mtprotoedge/outbound.go +++ b/internal/mtprotoedge/outbound.go @@ -4,13 +4,15 @@ import ( "bufio" "context" "crypto/aes" + "crypto/cipher" + "encoding/binary" "errors" "fmt" "io" + "sync" + "sync/atomic" "time" - "github.com/gotd/ige" - "github.com/gotd/td/bin" "github.com/gotd/td/crypto" "github.com/gotd/td/mt" @@ -24,16 +26,30 @@ var ( ErrConnClosed = errors.New("mtproto connection closed") // ErrOutboundQueueFull 表示 best-effort update push 未能在预算内进入出站队列。 ErrOutboundQueueFull = errors.New("mtproto outbound queue full") + // ErrOutboundTrackedBudget 表示消息无法进入 Server 级 resend tracking 预算。可靠 RPC + // 响应会终止连接让客户端重试；best-effort update 仅丢弃加速推送，由 difference 恢复。 + ErrOutboundTrackedBudget = errors.New("mtproto outbound tracked byte budget exhausted") + ErrOutboundMessageTooLarge = errors.New("mtproto outbound message exceeds transport frame limit") ) const ( - maxOutboundQueue = 1024 + // 原实现为每条 Conn eager 分配 1024 + 256 个 outboundOp 槽；在大连接数下仅空队列 + // backing 就占据大量常驻内存。默认缩至 128 + 32，仍覆盖 TDesktop 启动/推送突发， + // 慢消费者继续由 best-effort timeout + durable difference 降级。 + defaultOutboundQueueSize = 128 + defaultOutboundControlQueueSize = 32 + defaultOutboundTrackedMaxBytes = int64(512 << 20) // 512 MiB / Server + defaultOutboundControlMaxBytes = int64(64 << 20) // ack/state/resend vectors / Server + maxTrackedServerMsgIDs = 4096 maxTrackedAckedMsgIDs = 1024 // maxTrackedServerBytes 是 pending（已发送待 ack、用于 resend）总 body 字节上限。 // 与 maxTrackedServerMsgIDs 并列：客户端从不 ack 时，大响应体按字节滚动丢弃， // 防 pending 被「4096 条 × 大 body」撑爆。 maxTrackedServerBytes = 64 << 20 // 64 MiB + // Encrypted transport adds auth-key/msg-key, plaintext headers, randomized + // padding and codec framing. Reject before creating the two encryption buffers. + maxOutboundBodyBytes = maxTransportMessageSize - (2 << 10) ) type outboundOpKind byte @@ -47,16 +63,21 @@ const ( ) type outboundOp struct { - kind outboundOpKind - control bool - ctx context.Context - msgType proto.MessageType - msg bin.Encoder - encoded *encodedOutboundMessage - ids []int64 - reqMsgID int64 - enqueuedAt time.Time - done chan outboundResult + kind outboundOpKind + control bool + ctx context.Context + msgType proto.MessageType + msg bin.Encoder + encoded *encodedOutboundMessage + // reservedBytes accounts for the encoded body while it is queued. For a + // content frame the reservation is transferred to resend tracking after a + // successful write; every other terminal path releases it exactly once. + reservedBytes int + reservationBudget *outboundTrackedBudget + ids []int64 + reqMsgID int64 + enqueuedAt time.Time + done chan outboundResult } type encodedOutboundMessage struct { @@ -72,29 +93,212 @@ type outboundResult struct { } type outboundFrame struct { - msgID int64 - seqNo int32 - typeID uint32 - body []byte - reqMsgID int64 - sentAt time.Time - sends int + msgID int64 + seqNo int32 + typeID uint32 + body []byte + reservedBytes int + // reservationBudget follows this frame from producer queue through resend tracking. + // Service frames such as new_session_created are content-related and therefore pending, + // but their bytes must remain on the independent control budget for the full lifetime. + reservationBudget *outboundTrackedBudget + reqMsgID int64 + sentAt time.Time + sends int } type outboundState struct { - pending map[int64]*outboundFrame - order []int64 - byRequest map[int64]int64 - acked map[int64]struct{} - ackOrder []int64 - totalBytes int + pending map[int64]*outboundFrame + order []int64 + byRequest map[int64]int64 + acked map[int64]struct{} + ackOrder []int64 + totalBytes int + maxMessages int + maxBytes int + budget *outboundTrackedBudget } -func newOutboundState() *outboundState { +// outboundTrackedBudget 是 body/control/write 三类预算共用的原子 byte-budget primitive。 +// 每个实例只负责一类：普通 encoded body、encoded service frame + 控制向量，或 write +// scratch；同一 reservation 不跨实例释放。 +// 预算在入队前预留，写成功后从 queue reservation 原子转交给 pending；CAS 避免所有 +// outbound producer/actor 在一把全局 mutex 上串行。 +type outboundTrackedBudget struct { + maxBytes int64 + used atomic.Int64 + wakeMu sync.Mutex + wake *outboundBudgetWake +} + +type outboundBudgetWake struct { + ch chan struct{} + waiters int +} + +const defaultOutboundEncodeConcurrency = 32 + +// outboundEncodeSlots bounds the otherwise unaccounted transient allocation made by TL +// encoding and layer transcoding. The retained encoded bytes are covered by +// outboundTrackedBudget after Encode returns, but without this process-wide gate many RPC +// workers could all allocate a near-limit body before any of them attempted that reservation. +// Encoding cannot be cancelled once an Encoder has started, so admission is intentionally +// acquired before calling user/domain supplied Encoder code and released immediately after the +// transient allocation has either become a tracked body or been discarded. +var outboundEncodeSlots = make(chan struct{}, defaultOutboundEncodeConcurrency) + +func withOutboundEncodeSlot(ctx context.Context, stop <-chan struct{}, fn func() error) error { + if ctx == nil { + ctx = context.Background() + } + select { + case outboundEncodeSlots <- struct{}{}: + case <-ctx.Done(): + return ctx.Err() + case <-stop: + return ErrConnClosed + } + defer func() { <-outboundEncodeSlots }() + return fn() +} + +func newOutboundTrackedBudget(maxBytes int64) *outboundTrackedBudget { + if maxBytes <= 0 { + maxBytes = defaultOutboundTrackedMaxBytes + } + return &outboundTrackedBudget{ + maxBytes: maxBytes, + wake: &outboundBudgetWake{ch: make(chan struct{})}, + } +} + +func (b *outboundTrackedBudget) reserve(n int) bool { + if n <= 0 { + return true + } + bytes := int64(n) + if b == nil || bytes > b.maxBytes { + return false + } + for { + used := b.used.Load() + if used > b.maxBytes-bytes { + return false + } + if b.used.CompareAndSwap(used, used+bytes) { + return true + } + } +} + +func (b *outboundTrackedBudget) release(n int) { + if b == nil || n <= 0 { + return + } + if used := b.used.Add(-int64(n)); used < 0 { + panic("mtprotoedge: outbound tracked byte budget released more than reserved") + } + // A release can make room for multiple independent writes. Broadcast to every waiter in the + // current generation; a capacity-1 token can strand all but one waiter even while bytes are + // available indefinitely. Generations are only allocated on the saturated slow path. + b.wakeMu.Lock() + if b.wake.waiters > 0 { + old := b.wake + b.wake = &outboundBudgetWake{ch: make(chan struct{})} + close(old.ch) + } + b.wakeMu.Unlock() +} + +func (b *outboundTrackedBudget) waitReserve(ctx context.Context, stop <-chan struct{}, n int) error { + return b.waitReserveUntil(ctx, stop, n, time.Time{}) +} + +// waitReserveUntil is the saturated-path variant used by write scratch. The absolute deadline +// is observed from the first capacity wait, not only after encryption when socket I/O begins. +// Its timer is allocated lazily after the CAS fast path fails, preserving the allocation-free +// steady state. +func (b *outboundTrackedBudget) waitReserveUntil(ctx context.Context, stop <-chan struct{}, n int, deadline time.Time) error { + if b == nil || n < 0 || int64(n) > b.maxBytes { + return ErrOutboundTrackedBudget + } + if ctx == nil { + ctx = context.Background() + } + if d, ok := ctx.Deadline(); ok && (deadline.IsZero() || d.Before(deadline)) { + deadline = d + } + var ( + deadlineTimer *time.Timer + deadlineC <-chan time.Time + ) + defer func() { + if deadlineTimer != nil { + deadlineTimer.Stop() + } + }() + for { + if b.reserve(n) { + return nil + } + if !deadline.IsZero() && !time.Now().Before(deadline) { + return context.DeadlineExceeded + } + // Subscribe under the same lock used by release, then retry once while holding it. This + // closes the release-between-check-and-subscribe window without putting the normal CAS + // reservation path behind a global mutex. + b.wakeMu.Lock() + if b.reserve(n) { + b.wakeMu.Unlock() + return nil + } + generation := b.wake + generation.waiters++ + b.wakeMu.Unlock() + + if deadlineC == nil && !deadline.IsZero() { + wait := time.Until(deadline) + if wait < 0 { + wait = 0 + } + deadlineTimer = time.NewTimer(wait) + deadlineC = deadlineTimer.C + } + var err error + select { + case <-generation.ch: + case <-ctx.Done(): + err = ctx.Err() + case <-deadlineC: + err = context.DeadlineExceeded + case <-stop: + err = ErrConnClosed + } + b.wakeMu.Lock() + generation.waiters-- + b.wakeMu.Unlock() + if err != nil { + return err + } + } +} + +func (b *outboundTrackedBudget) snapshot() int64 { + if b == nil { + return 0 + } + return b.used.Load() +} + +func newOutboundState(budget *outboundTrackedBudget) *outboundState { + return newOutboundStateWithLimits(budget, maxTrackedServerMsgIDs, maxTrackedServerBytes) +} + +func newOutboundStateWithLimits(budget *outboundTrackedBudget, maxMessages, maxBytes int) *outboundState { return &outboundState{ - pending: make(map[int64]*outboundFrame), - byRequest: make(map[int64]int64), - acked: make(map[int64]struct{}), + maxMessages: maxMessages, + maxBytes: maxBytes, + budget: budget, } } @@ -102,8 +306,15 @@ func (c *Conn) startOutbound() { if c.metrics == nil { c.metrics = NopMetrics{} } - c.outbound = make(chan outboundOp, maxOutboundQueue) - c.outboundControl = make(chan outboundOp, maxOutboundQueue/4) + if c.outboundQueueSize <= 0 { + c.outboundQueueSize = defaultOutboundQueueSize + } + if c.outboundControlQueueSize <= 0 { + c.outboundControlQueueSize = defaultOutboundControlQueueSize + } + c.ensureOutboundTrackedBudget() + c.outbound = make(chan outboundOp, c.outboundQueueSize) + c.outboundControl = make(chan outboundOp, c.outboundControlQueueSize) c.outboundStop = make(chan struct{}) c.outboundDone = make(chan struct{}) go c.outboundLoop() @@ -111,23 +322,74 @@ func (c *Conn) startOutbound() { // Close 停止连接的出站 actor。它不关闭底层 transport；transport 生命周期仍由 serveConn 管理。 func (c *Conn) Close() { + c.beginTerminalShutdown() c.closeInboundRPCScheduler() - c.outboundClose.Do(func() { - if c.outboundStop != nil { - close(c.outboundStop) - <-c.outboundDone - } - }) + c.waitOutboundShutdown() +} + +// beginTerminalShutdown is the non-blocking ownership transition shared by graceful and hard +// close. It closes both producer gates and cancels RPC work before any potentially blocking +// transport.Close call, so a timed-out batch close cannot keep accepting memory/work. +func (c *Conn) beginTerminalShutdown() { + c.terminal.Store(true) + c.signalOutboundStop() + c.beginCloseInboundRPCScheduler() +} + +func (c *Conn) waitOutboundShutdown() { + if c.outboundDone != nil { + <-c.outboundDone + } } // ForceClose 停止连接并关闭底层 transport。 // 仅用于授权撤销 / destroy_auth_key 这类“必须让对端立即断线”的路径；普通生命周期仍由 // serveConn 统一关闭 transport，避免正常 push/索引清理把长连接误伤成硬断。 func (c *Conn) ForceClose() { - if c.transport != nil { - _ = c.transport.Close() - } - c.Close() + c.beginTerminalShutdown() + c.closeTransport() + c.closeInboundRPCScheduler() + c.waitOutboundShutdown() +} + +// closeTransport 只关闭物理 transport，不等待 outbound actor。写失败路径运行在 +// actor 自身 goroutine 中，若在这里调用 Close 会等待 outboundDone 而自锁。 +func (c *Conn) closeTransport() { + c.transportClose.Do(func() { + if c.transport != nil { + _ = c.transport.Close() + } + }) +} + +// failTransport 把不可恢复的写错误提升为连接级 terminal failure。它只负责 +// 标记 terminal + 关闭 socket；handleOutboundOp 返回后，actor 自己发停止信号并退出， +// serveConn 被 Close 解开 Recv 后负责注销索引。 +func (c *Conn) failTransport() { + // Publish both producer gates before Close: a custom/broken transport may block in + // Close, but it must not keep accepting queued bodies or RPC work meanwhile. + c.beginTerminalShutdown() + c.closeTransport() +} + +// dropSlowConsumer 把出站队列持续拥塞的连接降级为离线连接。它不能等待 outbound +// actor：调用方位于 fan-out 热路径，等待单个慢 socket 会把同一用户的健康设备和 +// transactional outbox lane 一起拖住。关闭 transport 会打断可能阻塞的写；serveConn +// 随后负责 Unregister，durable update 由该设备重连后的 getDifference 补偿。 +func (c *Conn) dropSlowConsumer() { + c.beginTerminalShutdown() + c.closeTransport() +} + +func (c *Conn) signalOutboundStop() { + c.outboundClose.Do(func() { + c.outboundEnqueueMu.Lock() + c.outboundClosing = true + if c.outboundStop != nil { + close(c.outboundStop) + } + c.outboundEnqueueMu.Unlock() + }) } // Send 加密并发送一条 server 消息。 @@ -154,28 +416,45 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin. if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed } + if c.terminal.Load() { + return ErrConnClosed + } writeCtx := context.Background() if ctx != nil { writeCtx = context.WithoutCancel(ctx) } - op := outboundOp{ - kind: outboundSend, - ctx: writeCtx, - msgType: t, - msg: msg, - encoded: encoded, - enqueuedAt: time.Now(), + // Encode before registering as an enqueue owner. Encoder is an external interface and may + // block forever; connection shutdown must not wait on it. The subsequent producer gate either + // accepts the completed/tracked body or rejects it and releases the reservation. + op, err := c.newOutboundSendOp(ctx, t, msg, encoded, false) + if err != nil { + // Best-effort durable updates may be dropped under process-wide pressure and + // recovered via getDifference. Do not close this healthy connection merely because + // other slow sockets currently own the shared body budget. + if errors.Is(err, ErrOutboundTrackedBudget) { + c.metrics.OutboundDropped("tracked_global_byte_budget") + } + return err } + if !c.beginOutboundEnqueue() { + op.releaseReservation(c.outboundTrackedBudget) + return ErrConnClosed + } + defer c.endOutboundEnqueue() + op.ctx = writeCtx + op.enqueuedAt = time.Now() // 快路径：非阻塞入队。fan-out 每 (conn × push) 都走这里，队列有空位时不为 // 本次推送分配任何 timer（此前 timeout>0 无条件 WithTimeout，稳态白建 timer）。 select { case c.outbound <- op: return nil case <-c.outboundStop: + op.releaseReservation(c.outboundTrackedBudget) return ErrConnClosed default: } if timeout == 0 { + op.releaseReservation(c.outboundTrackedBudget) c.metrics.OutboundDropped("push_queue_full") return ErrOutboundQueueFull } @@ -193,11 +472,14 @@ func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin. case c.outbound <- op: return nil case <-timeoutC: + op.releaseReservation(c.outboundTrackedBudget) c.metrics.OutboundDropped("push_queue_timeout") return ErrOutboundQueueFull case <-ctx.Done(): + op.releaseReservation(c.outboundTrackedBudget) return ctx.Err() case <-c.outboundStop: + op.releaseReservation(c.outboundTrackedBudget) return ErrConnClosed } } @@ -214,19 +496,25 @@ func (c *Conn) sendOutbound(ctx context.Context, t proto.MessageType, msg bin.En if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed } - op := outboundOp{ - kind: outboundSend, - control: control, - ctx: ctx, - msgType: t, - msg: msg, - encoded: encoded, - enqueuedAt: time.Now(), - done: make(chan outboundResult, 1), - } - if err := c.enqueueOutbound(ctx, op); err != nil { + op, err := c.newOutboundSendOp(ctx, t, msg, encoded, control) + if err != nil { + c.failOutboundBudget(err) return err } + if !c.beginOutboundEnqueue() { + op.releaseReservation(c.outboundTrackedBudget) + return ErrConnClosed + } + op.control = control + op.ctx = ctx + op.enqueuedAt = time.Now() + op.done = make(chan outboundResult, 1) + if err := c.enqueueOutboundRegistered(ctx, op); err != nil { + op.releaseReservation(c.outboundTrackedBudget) + c.endOutboundEnqueue() + return err + } + c.endOutboundEnqueue() select { case res := <-op.done: return res.err @@ -245,21 +533,31 @@ func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encod if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed } - op := outboundOp{ - kind: outboundSend, - control: true, - ctx: ctx, - msgType: t, - msg: msg, - enqueuedAt: time.Now(), - // done 为 nil：fire-and-forget，handleOutboundSend 的 finish 对 nil done 安全跳过。 + if c.terminal.Load() { + return ErrConnClosed } + op, err := c.newOutboundSendOp(ctx, t, msg, nil, true) + if err != nil { + c.failOutboundBudget(err) + return err + } + if !c.beginOutboundEnqueue() { + op.releaseReservation(c.outboundTrackedBudget) + return ErrConnClosed + } + defer c.endOutboundEnqueue() + op.control = true + op.ctx = ctx + op.enqueuedAt = time.Now() + // done 为 nil：fire-and-forget，handleOutboundSend 的 finish 对 nil done 安全跳过。 select { case c.outboundControl <- op: return nil case <-c.outboundStop: + op.releaseReservation(c.outboundTrackedBudget) return ErrConnClosed default: + op.releaseReservation(c.outboundTrackedBudget) c.metrics.OutboundDropped("control_queue_full") return nil } @@ -267,15 +565,25 @@ func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encod // AckServerMessages 接收客户端 msgs_ack，释放已确认的 server 出站消息。 func (c *Conn) AckServerMessages(ids []int64) { - if len(ids) == 0 || c.outbound == nil || c.outboundControl == nil { + if len(ids) == 0 || c.outbound == nil || c.outboundControl == nil || c.terminal.Load() { return } - copied := append([]int64(nil), ids...) - op := outboundOp{kind: outboundAck, control: true, ids: copied} + op, err := c.newOutboundVectorOp(outboundAck, ids) + if err != nil { + c.failOutboundBudget(err) + return + } + if !c.beginOutboundEnqueue() { + op.releaseReservation(c.outboundTrackedBudget) + return + } + defer c.endOutboundEnqueue() select { case c.outboundControl <- op: case <-c.outboundStop: + op.releaseReservation(c.outboundTrackedBudget) default: + op.releaseReservation(c.outboundTrackedBudget) c.metrics.OutboundDropped("ack_queue_full") } } @@ -286,14 +594,15 @@ func (c *Conn) OutgoingStateInfo(ctx context.Context, ids []int64) ([]byte, erro if c.outbound == nil { return nil, ErrConnClosed } - op := outboundOp{ - kind: outboundQueryState, - control: true, - ctx: ctx, - ids: append([]int64(nil), ids...), - done: make(chan outboundResult, 1), + op, err := c.newOutboundVectorOp(outboundQueryState, ids) + if err != nil { + c.failOutboundBudget(err) + return nil, err } + op.ctx = ctx + op.done = make(chan outboundResult, 1) if err := c.enqueueOutbound(ctx, op); err != nil { + op.releaseReservation(c.outboundTrackedBudget) return nil, err } select { @@ -311,14 +620,15 @@ func (c *Conn) ResendMessages(ctx context.Context, ids []int64) ([]byte, error) if c.outbound == nil { return nil, ErrConnClosed } - op := outboundOp{ - kind: outboundResend, - control: true, - ctx: ctx, - ids: append([]int64(nil), ids...), - done: make(chan outboundResult, 1), + op, err := c.newOutboundVectorOp(outboundResend, ids) + if err != nil { + c.failOutboundBudget(err) + return nil, err } + op.ctx = ctx + op.done = make(chan outboundResult, 1) if err := c.enqueueOutbound(ctx, op); err != nil { + op.releaseReservation(c.outboundTrackedBudget) return nil, err } select { @@ -357,9 +667,20 @@ func (c *Conn) ResendByRequest(ctx context.Context, reqMsgID int64) (bool, error } func (c *Conn) enqueueOutbound(ctx context.Context, op outboundOp) error { + if !c.beginOutboundEnqueue() { + return ErrConnClosed + } + defer c.endOutboundEnqueue() + return c.enqueueOutboundRegistered(ctx, op) +} + +func (c *Conn) enqueueOutboundRegistered(ctx context.Context, op outboundOp) error { if ctx == nil { ctx = context.Background() } + if c.terminal.Load() { + return ErrConnClosed + } q := c.outbound if op.control { q = c.outboundControl @@ -384,13 +705,49 @@ func (c *Conn) enqueueOutbound(ctx context.Context, op outboundOp) error { } } +func (c *Conn) beginOutboundEnqueue() bool { + c.outboundEnqueueMu.Lock() + defer c.outboundEnqueueMu.Unlock() + if c.outboundClosing || c.terminal.Load() { + return false + } + c.outboundEnqueueWG.Add(1) + return true +} + +func (c *Conn) endOutboundEnqueue() { + c.outboundEnqueueWG.Done() +} + func (c *Conn) outboundLoop() { - defer close(c.outboundDone) - state := newOutboundState() + state := newOutboundState(c.outboundTrackedBudget) + defer func() { + // pending frames belong exclusively to this actor. Releasing after drain ensures no + // resend path can race the final budget return and no Conn body survives actor exit. + state.releaseAll() + close(c.outboundDone) + }() for { + if c.terminal.Load() { + c.signalOutboundStop() + c.drainOutbound() + return + } select { case op := <-c.outboundControl: + if c.terminal.Load() { + op.releaseReservation(c.outboundTrackedBudget) + op.finish(outboundResult{err: ErrConnClosed}) + c.signalOutboundStop() + c.drainOutbound() + return + } c.handleOutboundOp(state, op) + if c.terminal.Load() { + c.signalOutboundStop() + c.drainOutbound() + return + } continue default: } @@ -399,19 +756,44 @@ func (c *Conn) outboundLoop() { c.drainOutbound() return case op := <-c.outboundControl: + if c.terminal.Load() { + op.releaseReservation(c.outboundTrackedBudget) + op.finish(outboundResult{err: ErrConnClosed}) + c.signalOutboundStop() + c.drainOutbound() + return + } c.handleOutboundOp(state, op) case op := <-c.outbound: + if c.terminal.Load() { + op.releaseReservation(c.outboundTrackedBudget) + op.finish(outboundResult{err: ErrConnClosed}) + c.signalOutboundStop() + c.drainOutbound() + return + } c.handleOutboundOp(state, op) } + if c.terminal.Load() { + c.signalOutboundStop() + c.drainOutbound() + return + } } } func (c *Conn) drainOutbound() { + // signalOutboundStop closes the producer gate before the actor gets here. + // Waiting first guarantees every producer either enqueued an owned op or + // released its reservation, after which this final drain cannot miss a body. + c.outboundEnqueueWG.Wait() for { select { case op := <-c.outboundControl: + op.releaseReservation(c.outboundTrackedBudget) op.finish(outboundResult{err: ErrConnClosed}) case op := <-c.outbound: + op.releaseReservation(c.outboundTrackedBudget) op.finish(outboundResult{err: ErrConnClosed}) default: return @@ -420,6 +802,9 @@ func (c *Conn) drainOutbound() { } func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) { + if op.kind != outboundSend { + defer op.releaseReservation(state.budget) + } switch op.kind { case outboundSend: c.handleOutboundSend(state, op) @@ -439,19 +824,53 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) { } func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) { - frame, err := c.buildFrame(op.msgType, op.msg, op.encoded) + frame, err := c.buildFrame(op.ctx, op.msgType, op.msg, op.encoded) + reserved := op.reservedBytes + reservationBudget := op.reservationBudget + if reservationBudget == nil { + reservationBudget = state.budget + } + op.reservedBytes = 0 + op.reservationBudget = nil + // A per-connection layer downgrade can allocate a different body. Reserve the + // replacement before dropping the canonical queue reservation so the transient + // two-body peak is also covered by the Server budget. + if err == nil && frame != nil && op.encoded != nil && !sameBacking(frame.body, op.encoded.body) { + if !reservationBudget.reserve(len(frame.body)) { + err = ErrOutboundTrackedBudget + } else { + reservationBudget.release(reserved) + reserved = len(frame.body) + op.encoded = nil + } + } + if err == nil && frame != nil && frameNeedsAck(frame.typeID) { + // The queue reservation is transferred to pending after write. A frame larger + // than the per-Conn resend ceiling is rejected before any bytes hit the wire. + if len(frame.body) > maxTrackedServerBytes { + err = ErrOutboundTrackedBudget + } + } + if errors.Is(err, ErrOutboundTrackedBudget) { + c.metrics.OutboundDropped("tracked_global_byte_budget") + c.failTransport() + } if err == nil { err = c.writeFrame(op.ctx, frame) } if err == nil && frame != nil && frameNeedsAck(frame.typeID) { // 写成功后才提交 content seq_no 递增（peekSeqNo 已按当前计数算好本帧 seq_no）。 c.commitContentSeqNo() - if dropped := state.add(frame); dropped > 0 { + frame.reservedBytes = reserved + frame.reservationBudget = reservationBudget + reserved = 0 + if dropped := state.addReserved(frame); dropped > 0 { for i := 0; i < dropped; i++ { c.metrics.OutboundDropped("tracked_queue_overflow") } } } + reservationBudget.release(reserved) queueWait := time.Since(op.enqueuedAt) bytes := 0 typeID := uint32(0) @@ -515,7 +934,134 @@ func (op outboundOp) finish(res outboundResult) { } } -func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage) (*outboundFrame, error) { +func (op *outboundOp) releaseReservation(budget *outboundTrackedBudget) { + if op == nil || op.reservedBytes <= 0 { + return + } + if op.reservationBudget != nil { + budget = op.reservationBudget + } + bytes := op.reservedBytes + op.reservedBytes = 0 + op.encoded = nil + op.msg = nil + op.ids = nil + op.reservationBudget = nil + // Make the queued body/vector unreachable before advertising its bytes to another producer. + budget.release(bytes) +} + +func (c *Conn) newOutboundVectorOp(kind outboundOpKind, ids []int64) (outboundOp, error) { + bytes := len(ids) * 8 + budget := c.ensureOutboundControlTrackedBudget() + if !budget.reserve(bytes) { + return outboundOp{}, ErrOutboundTrackedBudget + } + return outboundOp{ + kind: kind, + control: true, + ids: append([]int64(nil), ids...), + reservedBytes: bytes, + reservationBudget: budget, + }, nil +} + +func (c *Conn) newOutboundSendOp(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage, priorityControl bool) (outboundOp, error) { + var budget *outboundTrackedBudget + if encoded == nil { + var bytes int + err := withOutboundEncodeSlot(ctx, c.outboundStop, func() error { + var err error + encoded, err = encodeOutboundMessageWithoutSlot(msg) + if err != nil { + return err + } + if encoded == nil { + return errors.New("nil encoded outbound message") + } + bytes = len(encoded.body) + if bytes > maxOutboundBodyBytes { + return fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, bytes, maxOutboundBodyBytes) + } + budget = c.outboundMessageBudget(encoded.typeID, priorityControl) + // Keep the transient encode slot until the completed body has entered the + // retained-byte budget. Otherwise goroutines could successively finish an + // encode, be descheduled before reserve, and accumulate an unbounded number + // of completed-but-untracked bodies despite the encode concurrency gate. + if !budget.reserve(bytes) { + return ErrOutboundTrackedBudget + } + return nil + }) + if err != nil { + return outboundOp{}, err + } + return outboundOp{ + kind: outboundSend, + msgType: t, + encoded: encoded, + reservedBytes: bytes, + reservationBudget: budget, + }, nil + } + if encoded == nil { + return outboundOp{}, errors.New("nil encoded outbound message") + } + bytes := len(encoded.body) + if bytes > maxOutboundBodyBytes { + return outboundOp{}, fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, bytes, maxOutboundBodyBytes) + } + budget = c.outboundMessageBudget(encoded.typeID, priorityControl) + if !budget.reserve(bytes) { + return outboundOp{}, ErrOutboundTrackedBudget + } + return outboundOp{ + kind: outboundSend, + msgType: t, + encoded: encoded, + reservedBytes: bytes, + reservationBudget: budget, + }, nil +} + +func (c *Conn) outboundMessageBudget(typeID uint32, priorityControl bool) *outboundTrackedBudget { + if priorityControl || encodedControlFrame(typeID) { + return c.ensureOutboundControlTrackedBudget() + } + return c.ensureOutboundTrackedBudget() +} + +func (c *Conn) failOutboundBudget(err error) { + if !errors.Is(err, ErrOutboundTrackedBudget) { + return + } + if c.metrics != nil { + c.metrics.OutboundDropped("tracked_global_byte_budget") + } + c.failTransport() +} + +func (c *Conn) ensureOutboundTrackedBudget() *outboundTrackedBudget { + c.outboundBudgetOnce.Do(func() { + if c.outboundTrackedBudget == nil { + // Standalone tests/embedders still get a bounded budget. Server-created + // connections receive the shared Server budget before this can run. + c.outboundTrackedBudget = newOutboundTrackedBudget(defaultOutboundTrackedMaxBytes) + } + }) + return c.outboundTrackedBudget +} + +func (c *Conn) ensureOutboundControlTrackedBudget() *outboundTrackedBudget { + c.outboundControlBudgetOnce.Do(func() { + if c.outboundControlTrackedBudget == nil { + c.outboundControlTrackedBudget = newOutboundTrackedBudget(defaultOutboundControlMaxBytes) + } + }) + return c.outboundControlTrackedBudget +} + +func (c *Conn) buildFrame(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage) (*outboundFrame, error) { if encoded == nil { var err error encoded, err = encodeOutboundMessage(msg) @@ -528,7 +1074,7 @@ func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder, encoded *encoded // 连接共享，故必须在此**逐连接**降级，且**绝不改共享 encoded**(downgradedClone 拷贝)。 // - rpc_result 的内层对象已在 encodeRPCResult 按 layer 降级，其 mt.* 外壳在此为顶层直通(no-op)。 // - 控制消息(mt.*)顶层直通。layer>=227 整条零开销。 - encoded = c.downgradedClone(encoded) + encoded = c.downgradedCloneContext(ctx, encoded) content := frameNeedsAck(encoded.typeID) msgID := c.msgID.New(t) return &outboundFrame{ @@ -545,13 +1091,22 @@ func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder, encoded *encoded // layer>=227 或 Transcode 直通(mt.* / 无变化)时原样返回入参，零拷贝。降级失败 fail-safe： // 返回 canonical 并计 metrics（宁可老客户端对个别长尾对象渲染异常，也不让连接/流崩）。 func (c *Conn) downgradedClone(encoded *encodedOutboundMessage) *encodedOutboundMessage { + return c.downgradedCloneContext(context.Background(), encoded) +} + +func (c *Conn) downgradedCloneContext(ctx context.Context, encoded *encodedOutboundMessage) *encodedOutboundMessage { if encoded == nil { return nil } if c.ClientLayer() >= layerwire.CanonicalLayer { return encoded } - down, err := layerwire.Transcode(encoded.body, c.ClientLayer()) + var down []byte + err := withOutboundEncodeSlot(ctx, c.outboundStop, func() error { + var err error + down, err = layerwire.Transcode(encoded.body, c.ClientLayer()) + return err + }) if err != nil { c.metrics.OutboundDropped("layerwire_downgrade_failed") return encoded @@ -573,6 +1128,20 @@ func sameBacking(a, b []byte) bool { } func encodeOutboundMessage(msg bin.Encoder) (*encodedOutboundMessage, error) { + return encodeOutboundMessageContext(context.Background(), msg) +} + +func encodeOutboundMessageContext(ctx context.Context, msg bin.Encoder) (*encodedOutboundMessage, error) { + var encoded *encodedOutboundMessage + err := withOutboundEncodeSlot(ctx, nil, func() error { + var err error + encoded, err = encodeOutboundMessageWithoutSlot(msg) + return err + }) + return encoded, err +} + +func encodeOutboundMessageWithoutSlot(msg bin.Encoder) (*encodedOutboundMessage, error) { if msg == nil { return nil, errors.New("nil outbound message") } @@ -618,22 +1187,32 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { if ctx == nil { ctx = context.Background() } - out, err := c.encryptOutboundFrame(frame) + // One absolute deadline covers the complete write attempt, including global scratch + // admission. Previously writeTimeout only started after scratch was acquired, so one blocked + // writer could make unrelated connections wait for their much longer RPC context deadline. + deadline := c.outboundWriteDeadline(ctx) + pool := c.ensureOutboundScratchPool() + scratch, err := pool.acquireUntil(ctx, c.outboundStop, encryptedOutboundWireLen(len(frame.body)), deadline) + if err != nil { + return fmt.Errorf("reserve outbound write scratch: %w", err) + } + defer pool.release(scratch) + out, err := c.encryptOutboundFrameInto(frame, &scratch.wire) if err != nil { return fmt.Errorf("encrypt: %w", err) } + // Capacity may become available at the deadline boundary, or encryption itself may consume + // the remaining budget. No socket bytes exist yet, so return a non-terminal timeout instead + // of calling the writer with an already-expired deadline and misclassifying it as a possibly + // partial write. + if err := prewriteDeadlineError(ctx, deadline); err != nil { + return fmt.Errorf("outbound deadline before write: %w", err) + } writer := c.writer if writer == nil { writer = c.transport } - var deadline time.Time - if c.writeTimeout > 0 { - deadline = time.Now().Add(c.writeTimeout) - } - if d, ok := ctx.Deadline(); ok && (deadline.IsZero() || d.Before(deadline)) { - deadline = d - } if dw, ok := writer.(deadlineOutboundWriter); ok { err = dw.SendDeadline(deadline, out) } else { @@ -647,6 +1226,9 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { cancel() } if err != nil { + // 任一 partial write / timeout 都可能破坏 MTProto 帧边界；该 socket + // 不可继续复用。这里只发 terminal 信号，不在 actor 内等待自身退出。 + c.failTransport() return fmt.Errorf("send: %w", err) } if frame.sentAt.IsZero() { @@ -656,45 +1238,112 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { return nil } -func (c *Conn) encryptOutboundFrame(frame *outboundFrame) (*bin.Buffer, error) { - plain := &c.outboundPlain - plain.Reset() - plain.PutLong(c.salt) - plain.PutLong(c.sessionID) - plain.PutLong(frame.msgID) - plain.PutInt32(frame.seqNo) - plain.PutInt32(int32(len(frame.body))) - plain.Put(frame.body) +func (c *Conn) outboundWriteDeadline(ctx context.Context) time.Time { + var deadline time.Time + if c.writeTimeout > 0 { + deadline = time.Now().Add(c.writeTimeout) + } + if ctx != nil { + if d, ok := ctx.Deadline(); ok && (deadline.IsZero() || d.Before(deadline)) { + deadline = d + } + } + return deadline +} - paddingOffset := plain.Len() - paddingLen := encryptedPaddingLen(paddingOffset) - growBinBufferLen(plain, paddingOffset+paddingLen) +func prewriteDeadlineError(ctx context.Context, deadline time.Time) error { + if ctx != nil { + if err := ctx.Err(); err != nil { + return err + } + } + if !deadline.IsZero() && !time.Now().Before(deadline) { + return context.DeadlineExceeded + } + return nil +} + +func (c *Conn) ensureOutboundScratchPool() *outboundScratchPool { + c.outboundScratchOnce.Do(func() { + if c.outboundScratchPool == nil { + c.outboundScratchPool = newOutboundScratchPool(defaultOutboundWriteMaxBytes) + } + }) + return c.outboundScratchPool +} + +func (c *Conn) encryptOutboundFrame(frame *outboundFrame) (*bin.Buffer, error) { + wire := &bin.Buffer{Buf: make([]byte, encryptedOutboundWireLen(len(frame.body)))} + return c.encryptOutboundFrameInto(frame, wire) +} + +func encryptedOutboundWireLen(bodyLen int) int { + plainWithoutPadding := encryptedFrameHeaderLen + bodyLen + return 24 + plainWithoutPadding + encryptedPaddingLen(plainWithoutPadding) +} + +func (c *Conn) encryptOutboundFrameInto(frame *outboundFrame, wire *bin.Buffer) (*bin.Buffer, error) { + if frame == nil || wire == nil { + return nil, errors.New("nil outbound frame scratch") + } + wireLen := encryptedOutboundWireLen(len(frame.body)) + ensureBinBufferLen(wire, wireLen) + plain := wire.Buf[24:] + binary.LittleEndian.PutUint64(plain[0:8], uint64(c.salt)) + binary.LittleEndian.PutUint64(plain[8:16], uint64(c.sessionID)) + binary.LittleEndian.PutUint64(plain[16:24], uint64(frame.msgID)) + binary.LittleEndian.PutUint32(plain[24:28], uint32(frame.seqNo)) + binary.LittleEndian.PutUint32(plain[28:32], uint32(len(frame.body))) + copy(plain[encryptedFrameHeaderLen:], frame.body) + + paddingOffset := encryptedFrameHeaderLen + len(frame.body) // padding 随机数走 per-Conn 缓冲读：每帧 12..1024 字节直读 crypto/rand 是一次 // getrandom syscall，缓冲后按 ~1KiB 批量取。只由 outbound actor 单 goroutine 访问， // 随机源本身不变（仍是 cipher 的 CSPRNG），只是预读。 if c.outboundRand == nil { c.outboundRand = bufio.NewReaderSize(c.cipher.Rand(), 1024) } - if _, err := io.ReadFull(c.outboundRand, plain.Buf[paddingOffset:]); err != nil { + if _, err := io.ReadFull(c.outboundRand, plain[paddingOffset:]); err != nil { return nil, err } - msgKey := crypto.MessageKey(c.key.Value, plain.Raw(), crypto.Server) + msgKey := crypto.MessageKey(c.key.Value, plain, crypto.Server) key, iv := crypto.Keys(c.key.Value, msgKey, crypto.Server) aesBlock, err := aes.NewCipher(key[:]) if err != nil { return nil, err } - wireLen := len(c.key.ID) + len(msgKey) + plain.Len() - wire := &c.outboundWire - ensureBinBufferLen(wire, wireLen) copy(wire.Buf[:len(c.key.ID)], c.key.ID[:]) copy(wire.Buf[len(c.key.ID):len(c.key.ID)+len(msgKey)], msgKey[:]) - ige.EncryptBlocks(aesBlock, iv[:], wire.Buf[len(c.key.ID)+len(msgKey):], plain.Raw()) + encryptIGEInPlace(aesBlock, iv[:], plain) return wire, nil } +func encryptIGEInPlace(block cipher.Block, iv, buf []byte) { + blockSize := block.BlockSize() + if blockSize != aes.BlockSize || len(iv) != 2*blockSize || len(buf)%blockSize != 0 { + panic("mtprotoedge: invalid in-place IGE dimensions") + } + previousCipher := iv[:blockSize] + var previousPlain [aes.BlockSize]byte + copy(previousPlain[:], iv[blockSize:]) + for offset := 0; offset < len(buf); offset += blockSize { + current := buf[offset : offset+blockSize] + var currentPlain [aes.BlockSize]byte + copy(currentPlain[:], current) + for i := range current { + current[i] ^= previousCipher[i] + } + block.Encrypt(current, current) + for i := range current { + current[i] ^= previousPlain[i] + } + previousCipher = current + previousPlain = currentPlain + } +} + func encryptedPaddingLen(l int) int { return 16 + (16 - (l % 16)) } @@ -707,16 +1356,6 @@ func ensureBinBufferLen(b *bin.Buffer, n int) { b.Buf = b.Buf[:n] } -func growBinBufferLen(b *bin.Buffer, n int) { - if cap(b.Buf) < n { - next := make([]byte, n) - copy(next, b.Buf) - b.Buf = next - return - } - b.Buf = b.Buf[:n] -} - func frameNeedsAck(typeID uint32) bool { switch typeID { case mt.MsgsAckTypeID, @@ -735,6 +1374,35 @@ func frameNeedsAck(typeID uint32) bool { } } +// encodedControlFrame identifies MTProto service responses independently from content-related +// sequencing. new_session_created and destroy_session_* are content-related (and therefore stay +// in resend tracking until ACK), but their small protocol-critical bodies must not compete with +// RPC results/updates for the general outbound body budget. +func encodedControlFrame(typeID uint32) bool { + switch typeID { + case mt.MsgsAckTypeID, + mt.PongTypeID, + mt.FutureSaltsTypeID, + mt.BadMsgNotificationTypeID, + mt.BadServerSaltTypeID, + mt.MsgsStateInfoTypeID, + mt.MsgsAllInfoTypeID, + mt.MsgDetailedInfoTypeID, + mt.MsgNewDetailedInfoTypeID, + mt.NewSessionCreatedTypeID, + mt.DestroySessionOkTypeID, + mt.DestroySessionNoneTypeID, + mt.RPCAnswerUnknownTypeID, + mt.RPCAnswerDroppedRunningTypeID, + mt.RPCAnswerDroppedTypeID, + destroyAuthKeyOkTypeID, + destroyAuthKeyFailTypeID: + return true + default: + return false + } +} + func outboundRequestMsgID(msg bin.Encoder) int64 { switch v := msg.(type) { case *proto.Result: @@ -744,11 +1412,22 @@ func outboundRequestMsgID(msg bin.Encoder) int64 { } } -func (s *outboundState) add(frame *outboundFrame) int { +// addReserved 接管调用方已经取得的全局 body 预算。pending 的每个元素恰好对应一份 +// reservation；后续只有 removePending/releaseAll 能归还。 +func (s *outboundState) addReserved(frame *outboundFrame) int { + if _, exists := s.pending[frame.msgID]; exists { + panic("mtprotoedge: duplicate outbound msg_id inserted into resend tracking") + } + if s.pending == nil { + s.pending = make(map[int64]*outboundFrame) + } s.pending[frame.msgID] = frame s.order = append(s.order, frame.msgID) s.totalBytes += len(frame.body) if frame.reqMsgID != 0 { + if s.byRequest == nil { + s.byRequest = make(map[int64]int64) + } s.byRequest[frame.reqMsgID] = frame.msgID } return s.shrinkPending() @@ -756,18 +1435,12 @@ func (s *outboundState) add(frame *outboundFrame) int { func (s *outboundState) ack(ids []int64) { for _, id := range ids { - frame, ok := s.pending[id] - if !ok { + if !s.removePending(id) { continue } - delete(s.pending, id) - s.totalBytes -= len(frame.body) - if frame.reqMsgID != 0 { - delete(s.byRequest, frame.reqMsgID) - } s.markAcked(id) } - if len(s.order) > maxTrackedServerMsgIDs*2 { + if len(s.order) > s.maxMessages*2 { s.compactOrder() } } @@ -794,6 +1467,9 @@ func (s *outboundState) markAcked(id int64) { if _, ok := s.acked[id]; ok { return } + if s.acked == nil { + s.acked = make(map[int64]struct{}) + } s.acked[id] = struct{}{} s.ackOrder = append(s.ackOrder, id) for len(s.ackOrder) > maxTrackedAckedMsgIDs { @@ -805,23 +1481,61 @@ func (s *outboundState) markAcked(id int64) { func (s *outboundState) shrinkPending() int { dropped := 0 - for (len(s.pending) > maxTrackedServerMsgIDs || s.totalBytes > maxTrackedServerBytes) && len(s.order) > 0 { + for (len(s.pending) > s.maxMessages || s.totalBytes > s.maxBytes) && len(s.order) > 0 { oldest := s.order[0] s.order = s.order[1:] - frame, ok := s.pending[oldest] - if !ok { + if !s.removePending(oldest) { continue } - delete(s.pending, oldest) - s.totalBytes -= len(frame.body) - if frame.reqMsgID != 0 { - delete(s.byRequest, frame.reqMsgID) - } dropped++ } return dropped } +func (s *outboundState) removePending(id int64) bool { + frame, ok := s.pending[id] + if !ok { + return false + } + delete(s.pending, id) + bytes := len(frame.body) + s.totalBytes -= bytes + if frame.reqMsgID != 0 { + if mapped, exists := s.byRequest[frame.reqMsgID]; exists && mapped == id { + delete(s.byRequest, frame.reqMsgID) + } + } + // Clear the body reference before making these bytes available to another connection. + frame.body = nil + frame.releaseReservation(s.budget) + return true +} + +func (s *outboundState) releaseAll() { + for _, frame := range s.pending { + frame.body = nil + frame.releaseReservation(s.budget) + } + s.pending = nil + s.order = nil + s.byRequest = nil + s.totalBytes = 0 +} + +func (f *outboundFrame) releaseReservation(defaultBudget *outboundTrackedBudget) { + if f == nil || f.reservedBytes <= 0 { + return + } + budget := f.reservationBudget + if budget == nil { + budget = defaultBudget + } + bytes := f.reservedBytes + f.reservedBytes = 0 + f.reservationBudget = nil + budget.release(bytes) +} + func (s *outboundState) compactOrder() { filtered := s.order[:0] for _, id := range s.order { diff --git a/internal/mtprotoedge/outbound_scratch.go b/internal/mtprotoedge/outbound_scratch.go new file mode 100644 index 00000000..a840bd38 --- /dev/null +++ b/internal/mtprotoedge/outbound_scratch.go @@ -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() +} diff --git a/internal/mtprotoedge/outbound_test.go b/internal/mtprotoedge/outbound_test.go index 95b6c0b6..798b4eba 100644 --- a/internal/mtprotoedge/outbound_test.go +++ b/internal/mtprotoedge/outbound_test.go @@ -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 diff --git a/internal/mtprotoedge/passkey_e2e_test.go b/internal/mtprotoedge/passkey_e2e_test.go index e73b3bea..006d7e67 100644 --- a/internal/mtprotoedge/passkey_e2e_test.go +++ b/internal/mtprotoedge/passkey_e2e_test.go @@ -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) diff --git a/internal/mtprotoedge/quick_ack_deadline_test.go b/internal/mtprotoedge/quick_ack_deadline_test.go new file mode 100644 index 00000000..41e79681 --- /dev/null +++ b/internal/mtprotoedge/quick_ack_deadline_test.go @@ -0,0 +1,70 @@ +package mtprotoedge + +import ( + "context" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/crypto" +) + +type quickAckDeadlineProbe struct { + requested bool + deadline time.Time + token uint32 +} + +func (p *quickAckDeadlineProbe) ConsumeQuickAckRequested() bool { + if !p.requested { + return false + } + p.requested = false + return true +} + +func (p *quickAckDeadlineProbe) SendQuickAck(ctx context.Context, token uint32) error { + p.deadline, _ = ctx.Deadline() + p.token = token + return nil +} + +func (p *quickAckDeadlineProbe) SendQuickAckDeadline(deadline time.Time, token uint32) error { + p.deadline = deadline + p.token = token + return nil +} + +func (*quickAckDeadlineProbe) Send(context.Context, *bin.Buffer) error { return nil } +func (*quickAckDeadlineProbe) Recv(context.Context, *bin.Buffer) error { return nil } +func (*quickAckDeadlineProbe) Close() error { return nil } + +func TestQuickAckUsesServerWriteDeadline(t *testing.T) { + probe := &quickAckDeadlineProbe{requested: true} + var key crypto.Key + authKey := key.WithID() + before := time.Now() + if err := sendQuickAckIfRequested(context.Background(), probe, authKey, []byte("plain"), 50*time.Millisecond); err != nil { + t.Fatalf("send quick ack: %v", err) + } + if probe.deadline.IsZero() { + t.Fatal("quick ack did not receive a write deadline") + } + if probe.deadline.Before(before.Add(40*time.Millisecond)) || probe.deadline.After(time.Now().Add(60*time.Millisecond)) { + t.Fatalf("quick ack deadline = %v, want about server timeout from now", probe.deadline) + } +} + +func TestQuickAckHonorsEarlierCallerDeadline(t *testing.T) { + probe := &quickAckDeadlineProbe{requested: true} + ctxDeadline := time.Now().Add(25 * time.Millisecond) + ctx, cancel := context.WithDeadline(context.Background(), ctxDeadline) + defer cancel() + var key crypto.Key + if err := sendQuickAckIfRequested(ctx, probe, key.WithID(), []byte("plain"), time.Second); err != nil { + t.Fatalf("send quick ack: %v", err) + } + if delta := probe.deadline.Sub(ctxDeadline); delta < -time.Millisecond || delta > time.Millisecond { + t.Fatalf("quick ack deadline = %v, want caller deadline %v", probe.deadline, ctxDeadline) + } +} diff --git a/internal/mtprotoedge/rpc_test.go b/internal/mtprotoedge/rpc_test.go index 8525dac2..8663134d 100644 --- a/internal/mtprotoedge/rpc_test.go +++ b/internal/mtprotoedge/rpc_test.go @@ -107,6 +107,118 @@ func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) { close(handler.release) } +func TestInboundRPCQueuedDeadlineReturnsRPCTimeout(t *testing.T) { + const dc = 2 + handler := &queueDeadlineRPC{ + firstStarted: make(chan struct{}), + releaseFirst: make(chan struct{}), + } + addr, pub, _ := startTestServer(t, Options{ + DC: dc, + RPC: handler, + RPCMaxInflight: 1, + RPCQueueSize: 2, + RPCTimeout: 60 * time.Millisecond, + RPCGlobalWorkers: 1, + }) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + clientMsgID := proto.NewMessageIDGen(time.Now) + firstReqID := clientMsgID.New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, firstReqID, 1, &tg.HelpGetConfigRequest{}) + select { + case <-handler.firstStarted: + case <-time.After(time.Second): + t.Fatal("timed out waiting for first rpc to start") + } + + secondReqID := clientMsgID.New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, secondReqID, 3, &tg.HelpGetConfigRequest{}) + // 第一条故意忽略 context，使第二条越过自身从入队起计算的 deadline 后才有机会出队。 + time.Sleep(120 * time.Millisecond) + close(handler.releaseFirst) + + result := readRPCResultForRequest(t, conn, cipher, auth.AuthKey, secondReqID) + var rpcErr mt.RPCError + if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil { + t.Fatalf("decode rpc timeout: %v", err) + } + if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "RPC_TIMEOUT" { + t.Fatalf("rpc_error = %d %q, want 500 RPC_TIMEOUT", rpcErr.ErrorCode, rpcErr.ErrorMessage) + } + if calls := handler.calls.Load(); calls != 1 { + t.Fatalf("handler calls = %d, want 1 (expired queued RPC must not dispatch)", calls) + } +} + +func TestInboundRPCRunningDeadlineReturnsExactlyOneTimeout(t *testing.T) { + for _, tc := range []struct { + name string + honorContext bool + }{ + {name: "handler_honors_context", honorContext: true}, + {name: "handler_temporarily_ignores_context", honorContext: false}, + } { + t.Run(tc.name, func(t *testing.T) { + const dc = 2 + handler := &runningDeadlineRPC{ + started: make(chan struct{}), + release: make(chan struct{}), + honorContext: tc.honorContext, + } + addr, pub, _ := startTestServer(t, Options{ + DC: dc, + RPC: handler, + RPCMaxInflight: 1, + RPCQueueSize: 1, + RPCTimeout: 60 * time.Millisecond, + RPCGlobalWorkers: 1, + }) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + clientMsgID := proto.NewMessageIDGen(time.Now) + reqID := clientMsgID.New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, reqID, 1, &tg.HelpGetConfigRequest{}) + select { + case <-handler.started: + case <-time.After(time.Second): + t.Fatal("timed out waiting for running rpc") + } + + // In the ignore-context case this result must arrive before release is closed: the + // scheduler deadline, not eventual handler return, owns the timeout response. + result := readRPCResultForRequest(t, conn, cipher, auth.AuthKey, reqID) + var rpcErr mt.RPCError + if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil { + t.Fatalf("decode running rpc timeout: %v", err) + } + if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "RPC_TIMEOUT" { + t.Fatalf("rpc_error = %d %q, want 500 RPC_TIMEOUT", rpcErr.ErrorCode, rpcErr.ErrorMessage) + } + close(handler.release) + }) + } +} + +func TestRPCResponseGateExactlyOnce(t *testing.T) { + for i := 0; i < 100; i++ { + gate := &rpcResponseGate{} + results := make(chan bool, 2) + go func() { results <- gate.tryNormal() }() + go func() { results <- gate.tryTimeout() }() + wins := 0 + if <-results { + wins++ + } + if <-results { + wins++ + } + if wins != 1 { + t.Fatalf("iteration %d response gate winners = %d, want 1", i, wins) + } + } +} + func TestDuplicateRPCResultAcrossReconnectUsesSessionCache(t *testing.T) { const dc = 2 handler := &countingConfigRPC{} @@ -208,6 +320,40 @@ func (h *blockingRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.B func (h *blockingRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } +type queueDeadlineRPC struct { + calls atomic.Int32 + firstStarted chan struct{} + releaseFirst chan struct{} +} + +func (h *queueDeadlineRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) { + if h.calls.Add(1) == 1 { + close(h.firstStarted) + <-h.releaseFirst + } + return &tg.Config{ThisDC: 2}, nil +} + +func (h *queueDeadlineRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + +type runningDeadlineRPC struct { + started chan struct{} + release chan struct{} + honorContext bool +} + +func (h *runningDeadlineRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) { + close(h.started) + if h.honorContext { + <-ctx.Done() + return nil, ctx.Err() + } + <-h.release + return &tg.Config{ThisDC: 2}, nil +} + +func (h *runningDeadlineRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + type canceledInternalRPC struct { calls atomic.Int32 firstDone chan struct{} diff --git a/internal/mtprotoedge/same_port_mux.go b/internal/mtprotoedge/same_port_mux.go index d4341bf0..72eff5b3 100644 --- a/internal/mtprotoedge/same_port_mux.go +++ b/internal/mtprotoedge/same_port_mux.go @@ -40,6 +40,13 @@ type samePortMux struct { closed chan struct{} once sync.Once + + // sniffing contains only sockets still owned by dispatch while it reads the first four + // bytes. Keeping an explicit registry lets Close interrupt every slow-loris read without a + // second cancellation goroutine per raw connection. A socket is removed under sniffMu before + // successful child-listener hand-off, establishing the ownership barrier. + sniffMu sync.Mutex + sniffing map[net.Conn]struct{} } func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux { @@ -51,6 +58,7 @@ func newSamePortMux(base net.Listener, sniffTimeout time.Duration) *samePortMux addr: base.Addr(), sniffTimeout: sniffTimeout, closed: make(chan struct{}), + sniffing: make(map[net.Conn]struct{}), } m.tcp = newSamePortMuxListener(m.addr, m.closed) m.http = newSamePortMuxListener(m.addr, m.closed) @@ -69,7 +77,15 @@ func (m *samePortMux) HTTP() net.Listener { func (m *samePortMux) Serve(ctx context.Context) error { ctx, cancel := context.WithCancel(ctx) - defer cancel() + // Every exit path must publish cancellation and close both child listeners before waiting + // for sniff/delivery goroutines. A permanent Accept error can otherwise leave a dispatch + // blocked on a full child backlog while the old defer order waits for it before canceling. + var wg sync.WaitGroup + defer func() { + cancel() + _ = m.Close() + wg.Wait() + }() go func() { <-ctx.Done() @@ -77,17 +93,23 @@ func (m *samePortMux) Serve(ctx context.Context) error { }() // 每条连接一个窥探 goroutine：wg 让 Serve 在退出前等待在途窥探把连接交接完成。 - var wg sync.WaitGroup - defer wg.Wait() - + var tempDelay time.Duration for { conn, err := m.base.Accept() if err != nil { if ctx.Err() != nil || isSamePortMuxClosed(m.closed) || isNetClosed(err) { return nil } + if isTemporaryAcceptError(err) { + tempDelay = nextAcceptRetryDelay(tempDelay) + if !waitAcceptRetry(ctx, tempDelay) { + return nil + } + continue + } return err } + tempDelay = 0 wg.Add(1) go func() { defer wg.Done() @@ -99,6 +121,19 @@ func (m *samePortMux) Serve(ctx context.Context) error { func (m *samePortMux) Close() error { m.once.Do(func() { close(m.closed) + // Snapshot under the ownership lock, then close outside it. finishSniff observes + // m.closed and refuses hand-off even after the map is cleared, so dispatch cannot race + // this snapshot and deliver a socket that Close is about to terminate. + m.sniffMu.Lock() + sniffing := make([]net.Conn, 0, len(m.sniffing)) + for conn := range m.sniffing { + sniffing = append(sniffing, conn) + delete(m.sniffing, conn) + } + m.sniffMu.Unlock() + for _, conn := range sniffing { + _ = conn.Close() + } _ = m.tcp.Close() _ = m.http.Close() _ = m.base.Close() @@ -109,6 +144,20 @@ func (m *samePortMux) Close() error { // dispatch 窥探单条连接的前 4 字节并把它交给 tcp 或 http 子 listener。窥探带 sniffTimeout // 读上界，慢/半开连接最多占用本 goroutine sniffTimeout 后即被回收。 func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) { + // SetReadDeadline bounds an otherwise healthy slow-loris connection, but Close only owns + // the base listener, not sockets Accept has already returned. Register temporary ownership so + // mux shutdown can close this read immediately. finishSniff removes the socket before hand-off. + if !m.beginSniff(conn) { + _ = conn.Close() + return + } + finishedSniff := false + defer func() { + if !finishedSniff { + m.finishSniff(conn) + } + }() + var header [4]byte if err := conn.SetReadDeadline(time.Now().Add(m.sniffTimeout)); err != nil { _ = conn.Close() @@ -118,6 +167,14 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) { _ = conn.Close() return } + // From this point onward deliver/child-listener closure owns cancellation. Removing the + // registry entry under sniffMu is the hand-off barrier: Close either captured and closed this + // socket, or it can no longer find it. A concurrently closed mux refuses delivery. + if !m.finishSniff(conn) { + _ = conn.Close() + return + } + finishedSniff = true if err := conn.SetReadDeadline(time.Time{}); err != nil { _ = conn.Close() return @@ -137,6 +194,30 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) { } } +func (m *samePortMux) beginSniff(conn net.Conn) bool { + m.sniffMu.Lock() + defer m.sniffMu.Unlock() + if isSamePortMuxClosed(m.closed) { + return false + } + if m.sniffing == nil { + m.sniffing = make(map[net.Conn]struct{}) + } + m.sniffing[conn] = struct{}{} + return true +} + +// finishSniff returns true only when dispatch still owned the socket and the mux remained open +// through the ownership barrier. A false result means Close captured the socket; dispatch must +// not hand it to a child listener. +func (m *samePortMux) finishSniff(conn net.Conn) bool { + m.sniffMu.Lock() + defer m.sniffMu.Unlock() + _, owned := m.sniffing[conn] + delete(m.sniffing, conn) + return owned && !isSamePortMuxClosed(m.closed) +} + // isHTTPHeaderPrefix 判断前 4 字节是否是 HTTP 请求行起始。 // // 这里只认 GET/POST/HEAD/OPTI，与 gotd generateInit 排除的前缀集合「严格对齐」：合法的 @@ -243,6 +324,10 @@ type samePortMuxListener struct { ch chan net.Conn closed chan struct{} once sync.Once + + deliveryMu sync.Mutex + closing bool + deliveryWG sync.WaitGroup } func newSamePortMuxListener(addr net.Addr, parentClosed <-chan struct{}) *samePortMuxListener { @@ -273,7 +358,25 @@ func (l *samePortMuxListener) Accept() (net.Conn, error) { func (l *samePortMuxListener) Close() error { l.once.Do(func() { + // Add and Wait on a WaitGroup must not race while the counter may still be zero. + // The delivery gate serializes the final Add with the transition to closing; after + // closing becomes true no producer can enter, so waiting and draining are safe. + l.deliveryMu.Lock() + l.closing = true close(l.closed) + l.deliveryMu.Unlock() + + l.deliveryWG.Wait() + for { + select { + case conn := <-l.ch: + if conn != nil { + _ = conn.Close() + } + default: + return + } + } }) return nil } @@ -283,6 +386,11 @@ func (l *samePortMuxListener) Addr() net.Addr { } func (l *samePortMuxListener) deliver(ctx context.Context, conn net.Conn) bool { + if !l.beginDelivery() { + return false + } + defer l.deliveryWG.Done() + select { case <-l.closed: return false @@ -292,3 +400,13 @@ func (l *samePortMuxListener) deliver(ctx context.Context, conn net.Conn) bool { return false } } + +func (l *samePortMuxListener) beginDelivery() bool { + l.deliveryMu.Lock() + defer l.deliveryMu.Unlock() + if l.closing { + return false + } + l.deliveryWG.Add(1) + return true +} diff --git a/internal/mtprotoedge/same_port_mux_test.go b/internal/mtprotoedge/same_port_mux_test.go new file mode 100644 index 00000000..8c40ef0a --- /dev/null +++ b/internal/mtprotoedge/same_port_mux_test.go @@ -0,0 +1,240 @@ +package mtprotoedge + +import ( + "bytes" + "context" + "errors" + "io" + "net" + "testing" + "time" +) + +func TestSamePortMuxListenerCloseWaitsAndReturnsBacklogAdmission(t *testing.T) { + admission := newAdmissionController(4, 4, 1) + listener := &samePortMuxListener{ + addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1)}, + ch: make(chan net.Conn, 1), + closed: make(chan struct{}), + } + + backlog, backlogPeer := trackedMuxPipe(t, admission, 1001) + defer backlogPeer.Close() + if !listener.deliver(context.Background(), backlog) { + t.Fatal("initial backlog delivery was rejected") + } + + // Deterministically model a producer that passed the delivery gate but has not yet + // completed. Close must publish closed first, then wait before draining the backlog. + if !listener.beginDelivery() { + t.Fatal("in-flight delivery gate unexpectedly closed") + } + closeDone := make(chan struct{}) + go func() { + _ = listener.Close() + close(closeDone) + }() + select { + case <-listener.closed: + case <-time.After(time.Second): + t.Fatal("Close did not publish listener closure") + } + select { + case <-closeDone: + t.Fatal("Close returned before in-flight delivery completed") + default: + } + listener.deliveryWG.Done() + select { + case <-closeDone: + case <-time.After(time.Second): + t.Fatal("Close did not finish after delivery completed") + } + + assertAdmissionConnections(t, admission, 0) + + late, latePeer := trackedMuxPipe(t, admission, 1002) + defer latePeer.Close() + if listener.deliver(context.Background(), late) { + t.Fatal("delivery after Close unexpectedly succeeded") + } + _ = late.Close() // dispatch owns and closes a rejected delivery. + assertAdmissionConnections(t, admission, 0) +} + +func TestSamePortMuxPermanentAcceptErrorCancelsBlockedDeliveryBeforeWait(t *testing.T) { + serverSide, clientSide := net.Pipe() + defer clientSide.Close() + wantErr := errors.New("same-port permanent accept failure") + base := &connThenErrorListener{conn: serverSide, err: wantErr} + closed := make(chan struct{}) + mux := &samePortMux{ + base: base, + addr: base.Addr(), + sniffTimeout: time.Hour, + closed: closed, + } + // An unbuffered child listener deterministically leaves dispatch blocked in deliver: no + // consumer is running, and the base listener immediately returns a permanent second error. + mux.tcp = &samePortMuxListener{addr: mux.addr, ch: make(chan net.Conn), closed: make(chan struct{})} + mux.http = &samePortMuxListener{addr: mux.addr, ch: make(chan net.Conn), closed: make(chan struct{})} + + writeDone := make(chan error, 1) + go func() { + _, err := clientSide.Write([]byte{0xef, 0, 0, 0}) + writeDone <- err + }() + serveDone := make(chan error, 1) + go func() { + serveDone <- mux.Serve(context.Background()) + }() + + select { + case err := <-serveDone: + if !errors.Is(err, wantErr) { + t.Fatalf("Serve error = %v, want %v", err, wantErr) + } + case <-time.After(time.Second): + t.Fatal("same-port Serve waited for blocked delivery before canceling it") + } + select { + case <-writeDone: + case <-time.After(time.Second): + t.Fatal("sniff writer remained blocked after same-port shutdown") + } +} + +func TestSamePortMuxShutdownInterruptsSlowSniffImmediately(t *testing.T) { + tests := []struct { + name string + shutdown func(context.CancelFunc, *samePortMux) + }{ + { + name: "context cancel", + shutdown: func(cancel context.CancelFunc, _ *samePortMux) { + cancel() + }, + }, + { + name: "mux close", + shutdown: func(_ context.CancelFunc, mux *samePortMux) { + _ = mux.Close() + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + base, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + mux := newSamePortMux(base, time.Minute) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + serveDone := make(chan error, 1) + go func() { serveDone <- mux.Serve(ctx) }() + + peer, err := net.Dial("tcp", base.Addr().String()) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer peer.Close() + // No bytes are written: dispatch is blocked in the four-byte sniff with a one-minute + // deadline. Shutdown must close this accepted socket instead of waiting for it. + tt.shutdown(cancel, mux) + + select { + case err := <-serveDone: + if err != nil { + t.Fatalf("Serve: %v", err) + } + case <-time.After(time.Second): + t.Fatal("Serve waited for the sniff deadline after shutdown") + } + if err := peer.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set peer deadline: %v", err) + } + var one [1]byte + if _, err := peer.Read(one[:]); err == nil { + t.Fatal("slow sniff socket remained open after mux shutdown") + } + }) + } +} + +func TestSamePortMuxSuccessfulHandoffReleasesSniffOwnership(t *testing.T) { + base, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + mux := newSamePortMux(base, time.Minute) + ctx, cancel := context.WithCancel(context.Background()) + serveDone := make(chan error, 1) + go func() { serveDone <- mux.Serve(ctx) }() + + peer, err := net.Dial("tcp", base.Addr().String()) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer peer.Close() + if _, err := peer.Write([]byte{0xef, 0, 0, 0}); err != nil { + t.Fatalf("write sniff prefix: %v", err) + } + accepted, err := mux.TCP().Accept() + if err != nil { + t.Fatalf("accept child: %v", err) + } + defer accepted.Close() + + // Once dispatch has delivered the Conn, canceling the mux may close listeners/backlog but + // must not let the old sniff watcher close a socket now owned by the child consumer. + cancel() + select { + case err := <-serveDone: + if err != nil { + t.Fatalf("Serve: %v", err) + } + case <-time.After(time.Second): + t.Fatal("Serve did not stop after cancel") + } + if _, err := peer.Write([]byte{1, 2, 3, 4}); err != nil { + t.Fatalf("write after handoff/shutdown: %v", err) + } + if err := accepted.SetReadDeadline(time.Now().Add(time.Second)); err != nil { + t.Fatalf("set accepted deadline: %v", err) + } + got := make([]byte, 8) + if _, err := io.ReadFull(accepted, got); err != nil { + t.Fatalf("read handed-off connection: %v", err) + } + want := []byte{0xef, 0, 0, 0, 1, 2, 3, 4} + if !bytes.Equal(got, want) { + t.Fatalf("handed-off bytes = %x, want %x", got, want) + } +} + +func trackedMuxPipe(t *testing.T, admission *admissionController, port int) (net.Conn, net.Conn) { + t.Helper() + server, peer := net.Pipe() + release, ok := admission.acquireConnection(&net.TCPAddr{ + IP: net.ParseIP("203.0.113.20"), + Port: port, + }) + if !ok { + _ = server.Close() + _ = peer.Close() + t.Fatal("test connection admission rejected") + } + return &admittedConn{Conn: server, release: release}, peer +} + +func assertAdmissionConnections(t *testing.T, admission *admissionController, want int) { + t.Helper() + admission.mu.Lock() + got := admission.connections + byIP := len(admission.byIP) + admission.mu.Unlock() + if got != want || (want == 0 && byIP != 0) { + t.Fatalf("admission state = connections:%d by_ip:%d, want connections:%d", got, byIP, want) + } +} diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index c1c22d9a..c9d09cd5 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -47,7 +47,9 @@ type RPCHandler interface { type Options struct { // Logger 日志器。默认 zap.NewNop()。 Logger *zap.Logger - // Codec 传输 codec 构造器。nil 表示自动探测（intermediate/abridged/full）。 + // Codec 传输 codec 构造器。nil 表示自动探测（intermediate/abridged/full）。自定义 + // codec 必须是 gotd 内置四种 codec（可包 NoHeader），或实现 InboundFrameBudgetedCodec； + // 无法在 payload 分配前预检长度的 codec 会 fail-closed。 Codec func() transport.Codec // ObfuscatedTCP 先按 MTProto TCP obfuscation 解包，再自动探测 codec。 // Telegram Desktop 的 tcpo_only endpoint 会走这个 64 字节前缀流程。 @@ -72,12 +74,45 @@ type Options struct { HandshakeMaxDuration time.Duration // WriteTimeout 单次写入超时。默认 30s。 WriteTimeout time.Duration + // MaxConnections 是进程接受的 raw 物理连接总上限，覆盖 codec sniff、握手和 + // 已认证连接的完整生命周期。默认 200000；负数表示不限制。 + MaxConnections int + // MaxConnectionsPerIP 是单 remote IP 的 raw 物理连接上限。默认 4096， + // 为共享 NAT 与 TDesktop 多候选连接保留足够突发；负数表示不限制。 + MaxConnectionsPerIP int + // MaxConcurrentHandshakes 是同时执行 auth_key_id=0 RSA/DH exchange 的上限。 + // 达限时已完成 transport framing 的连接收到 -429 后断开。默认 256；负数表示不限制。 + MaxConcurrentHandshakes int // RPCMaxInflight 是单连接同时处理的 RPC 上限。默认 32。 RPCMaxInflight int - // RPCQueueSize 是单连接等待处理的 RPC 队列长度。默认 256。 + // RPCQueueSize 是单连接等待处理的 RPC 队列长度。默认 64；队列按首条请求懒分配。 RPCQueueSize int // RPCTimeout 是单个 RPC 在连接层的最大处理时长。默认 30s。 + // 超时从 Copy 前预算/入队开始计算，排队时间包含在内。 RPCTimeout time.Duration + // RPCGlobalWorkers 是 Server 共享 inbound RPC worker 数。默认 256。 + RPCGlobalWorkers int + // RPCGlobalMaxTasks 是全进程已预留、排队和执行中的 RPC 条数上限。默认 8192。 + RPCGlobalMaxTasks int + // RPCGlobalMaxBytes 是上述 RPC body 的总字节预算。默认 512 MiB。 + RPCGlobalMaxBytes int64 + // InboundFrameGlobalMaxBytes 是所有物理连接当前正在处理的 transport wire buffer + // 与最大解密 plaintext buffer 的总预算。长度前缀读取后、payload 分配前预留，默认 + // 512 MiB；非正值使用默认值。 + InboundFrameGlobalMaxBytes int64 + // OutboundQueueSize / OutboundControlQueueSize 是每连接普通与控制 mailbox 容量。 + // 默认 128/32；控制队列在 actor 中保持严格优先。 + OutboundQueueSize int + OutboundControlQueueSize int + // OutboundTrackedGlobalMaxBytes 是所有连接为 msg_resend_req 保留的 RPC/update body + // 总预算。默认 512 MiB；编码后的 MTProto service frame 与控制向量另用 64 MiB + // control budget（包括需 resend tracking 的 new_session_created 等），避免 body 压力 + // 阻断连接维持消息。可靠响应无法 tracking 时终止该连接，durable best-effort update + // 则只丢在线加速并由 difference 恢复。 + OutboundTrackedGlobalMaxBytes int64 + // OutboundWriteGlobalMaxBytes bounds concurrent encrypted wire/codec/obfuscation scratch. + // Scratch is shared and pooled across connections; default 512 MiB. + OutboundWriteGlobalMaxBytes int64 // DC 是本 server 的 DC ID。默认 2。 DC int @@ -115,15 +150,48 @@ func (o *Options) setDefaults() { if o.WriteTimeout == 0 { o.WriteTimeout = 30 * time.Second } + if o.MaxConnections == 0 { + o.MaxConnections = defaultMaxConnections + } + if o.MaxConnectionsPerIP == 0 { + o.MaxConnectionsPerIP = defaultMaxConnectionsPerIP + } + if o.MaxConcurrentHandshakes == 0 { + o.MaxConcurrentHandshakes = defaultMaxConcurrentHandshakes + } if o.RPCMaxInflight <= 0 { o.RPCMaxInflight = 32 } if o.RPCQueueSize <= 0 { - o.RPCQueueSize = 256 + o.RPCQueueSize = 64 } if o.RPCTimeout == 0 { o.RPCTimeout = 30 * time.Second } + if o.RPCGlobalWorkers <= 0 { + o.RPCGlobalWorkers = 256 + } + if o.RPCGlobalMaxTasks <= 0 { + o.RPCGlobalMaxTasks = 8192 + } + if o.RPCGlobalMaxBytes <= 0 { + o.RPCGlobalMaxBytes = 512 << 20 + } + if o.InboundFrameGlobalMaxBytes <= 0 { + o.InboundFrameGlobalMaxBytes = defaultInboundFrameGlobalMaxBytes + } + if o.OutboundQueueSize <= 0 { + o.OutboundQueueSize = defaultOutboundQueueSize + } + if o.OutboundControlQueueSize <= 0 { + o.OutboundControlQueueSize = defaultOutboundControlQueueSize + } + if o.OutboundTrackedGlobalMaxBytes <= 0 { + o.OutboundTrackedGlobalMaxBytes = defaultOutboundTrackedMaxBytes + } + if o.OutboundWriteGlobalMaxBytes <= 0 { + o.OutboundWriteGlobalMaxBytes = defaultOutboundWriteMaxBytes + } if o.DC == 0 { o.DC = 2 } @@ -150,30 +218,38 @@ func (o *Options) setDefaults() { // 接受连接、协商 codec、完成密钥交换、解密并分发加密消息到 RPC 路由，处理服务消息， // 并把活跃连接注册到 SessionManager 以支持主动推送（updates 等）。不含业务逻辑。 type Server struct { - log *zap.Logger - codec func() transport.Codec - obfuscated bool - websocket bool - websocketOrigins []string - readTimeout time.Duration - handshakeTimeout time.Duration - handshakeMaxDur time.Duration - writeTimeout time.Duration - rpcInflight int - rpcQueueSize int - rpcTimeout time.Duration + log *zap.Logger + codec func() transport.Codec + obfuscated bool + websocket bool + websocketOrigins []string + readTimeout time.Duration + handshakeTimeout time.Duration + handshakeMaxDur time.Duration + writeTimeout time.Duration + rpcInflight int + rpcQueueSize int + rpcTimeout time.Duration + rpcScheduler *inboundRPCScheduler + frameBudget *inboundFrameBudget + outboundQueueSize int + outboundControlQueueSize int + outboundTrackedBudget *outboundTrackedBudget + outboundControlBudget *outboundTrackedBudget + outboundScratchPool *outboundScratchPool - dc int - key exchange.PrivateKey - authKeys store.AuthKeyStore - sessions store.SessionStore - conns *SessionManager - rpc RPCHandler - metrics Metrics - cipher crypto.Cipher - clock clock.Clock - rand io.Reader - types *tmap.Map + dc int + key exchange.PrivateKey + authKeys store.AuthKeyStore + sessions store.SessionStore + conns *SessionManager + rpc RPCHandler + metrics Metrics + cipher crypto.Cipher + clock clock.Clock + rand io.Reader + types *tmap.Map + admission *admissionController rpcResults *rpcResultCache @@ -189,30 +265,38 @@ func New(opts Options) *Server { conns = NewSessionManager(opts.Logger.Named("sessions")) } return &Server{ - log: opts.Logger, - codec: opts.Codec, - obfuscated: opts.ObfuscatedTCP, - websocket: opts.WebSocket, - websocketOrigins: append([]string(nil), opts.WebSocketAllowedOrigins...), - readTimeout: opts.ReadTimeout, - handshakeTimeout: opts.HandshakeIdleTimeout, - handshakeMaxDur: opts.HandshakeMaxDuration, - writeTimeout: opts.WriteTimeout, - rpcInflight: opts.RPCMaxInflight, - rpcQueueSize: opts.RPCQueueSize, - rpcTimeout: opts.RPCTimeout, - dc: opts.DC, - key: exchange.PrivateKey{RSA: opts.RSAKey}, - authKeys: opts.AuthKeys, - sessions: opts.Sessions, - conns: conns, - rpc: opts.RPC, - metrics: opts.Metrics, - cipher: crypto.NewServerCipher(opts.Rand), - clock: opts.Clock, - rand: opts.Rand, - types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()), - rpcResults: newRPCResultCache(opts.Clock.Now), + log: opts.Logger, + codec: opts.Codec, + obfuscated: opts.ObfuscatedTCP, + websocket: opts.WebSocket, + websocketOrigins: append([]string(nil), opts.WebSocketAllowedOrigins...), + readTimeout: opts.ReadTimeout, + handshakeTimeout: opts.HandshakeIdleTimeout, + handshakeMaxDur: opts.HandshakeMaxDuration, + writeTimeout: opts.WriteTimeout, + rpcInflight: opts.RPCMaxInflight, + rpcQueueSize: opts.RPCQueueSize, + rpcTimeout: opts.RPCTimeout, + rpcScheduler: newInboundRPCScheduler(opts.RPCGlobalWorkers, opts.RPCGlobalMaxTasks, opts.RPCGlobalMaxBytes), + frameBudget: newInboundFrameBudget(opts.InboundFrameGlobalMaxBytes), + outboundQueueSize: opts.OutboundQueueSize, + outboundControlQueueSize: opts.OutboundControlQueueSize, + outboundTrackedBudget: newOutboundTrackedBudget(opts.OutboundTrackedGlobalMaxBytes), + outboundControlBudget: newOutboundTrackedBudget(defaultOutboundControlMaxBytes), + outboundScratchPool: newOutboundScratchPool(opts.OutboundWriteGlobalMaxBytes), + dc: opts.DC, + key: exchange.PrivateKey{RSA: opts.RSAKey}, + authKeys: opts.AuthKeys, + sessions: opts.Sessions, + conns: conns, + rpc: opts.RPC, + metrics: opts.Metrics, + cipher: crypto.NewServerCipher(opts.Rand), + clock: opts.Clock, + rand: opts.Rand, + types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()), + rpcResults: newRPCResultCache(opts.Clock.Now), + admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes), } } @@ -224,27 +308,40 @@ func (s *Server) Conns() *SessionManager { // newConn 基于一次解密结果创建一个可发送的连接对象。 func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn { c := &Conn{ - transport: tc, - writer: tc, - cipher: s.cipher, - msgID: proto.NewMessageIDGen(s.clock.Now), - writeTimeout: s.writeTimeout, - metrics: s.metrics, - authKeyID: key.ID, - authKeyHex: hex.EncodeToString(key.ID[:]), - sessionID: sessionID, - salt: salt, - key: key, - createdAt: s.clock.Now(), + transport: tc, + writer: tc, + cipher: s.cipher, + msgID: proto.NewMessageIDGen(s.clock.Now), + writeTimeout: s.writeTimeout, + metrics: s.metrics, + authKeyID: key.ID, + authKeyHex: hex.EncodeToString(key.ID[:]), + sessionID: sessionID, + salt: salt, + key: key, + createdAt: s.clock.Now(), + outboundQueueSize: s.outboundQueueSize, + outboundControlQueueSize: s.outboundControlQueueSize, + outboundTrackedBudget: s.outboundTrackedBudget, + outboundControlTrackedBudget: s.outboundControlBudget, + outboundScratchPool: s.outboundScratchPool, } c.startOutbound() - c.startInboundRPCScheduler(s.rpcInflight, s.rpcQueueSize, s.rpcTimeout) + c.startInboundRPCScheduler(s.rpcScheduler, s.rpcInflight, s.rpcQueueSize, s.rpcTimeout) return c } // Serve 在 ln 上运行 MTProto 连接循环，直到 ctx 取消或发生不可恢复错误。 // ctx 取消时优雅退出：关闭 listener 并等待在途连接处理结束。 func (s *Server) Serve(ctx context.Context, ln net.Listener) error { + // 共享 worker 池只在 Server 真正 Serve 后允许消费，并在首条 RPC 到达时懒启动。 + // serveTCP/serveMixed 返回前会等待连接 goroutine 收敛，各 Conn 已先排空/取消任务； + // 最后再停止全局池，避免关闭过程中留下无人消费但仍占预算的队列。 + s.rpcScheduler.start() + defer s.rpcScheduler.stop(rpcCloseWaitTimeout) + // 只在最外层 listener 包一次，确保 same-port mux 的 sniff/HTTP upgrade 也计入 + // raw admission，而不是等连接已经分流后才计数。 + ln = s.admission.wrapListener(ln) if s.websocket { return s.serveMixed(ctx, ln) } @@ -292,11 +389,15 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { ) defer s.log.Info("Stopped") - go func() { - <-ctx.Done() + stopAll := func() { + cancel() _ = mux.Close() _ = httpServer.Close() _ = wsLn.Close() + } + go func() { + <-ctx.Done() + stopAll() }() errCh := make(chan error, 4) @@ -330,17 +431,19 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { errCh <- nil }() + // The four services form one lifecycle: even a clean/closed-listener return from any one + // component means the remaining three can no longer make forward progress as a complete + // same-port server. Stop them immediately, then collect their terminal results. var firstErr error - for i := 0; i < 4; i++ { + if err := <-errCh; err != nil { + firstErr = err + } + stopAll() + for i := 1; i < 4; i++ { if err := <-errCh; err != nil && firstErr == nil { firstErr = err - cancel() } } - cancel() - _ = mux.Close() - _ = httpServer.Close() - _ = wsLn.Close() wg.Wait() return firstErr } @@ -351,22 +454,39 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { // 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆（裸 MTProto TCP）；WebSocket // 连接传 false（gotd 升级处理器已完成去混淆）。 func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated bool) error { + ctx, cancel := context.WithCancel(ctx) + var wg sync.WaitGroup + defer func() { + // A permanent Accept error is itself a terminal lifecycle event. Cancel accepted + // connections and close the listener before waiting; otherwise a live connection can + // keep the WaitGroup blocked forever and prevent the accept error from being returned. + cancel() + _ = ln.Close() + wg.Wait() + }() go func() { <-ctx.Done() _ = ln.Close() }() - var wg sync.WaitGroup - defer wg.Wait() - + var tempDelay time.Duration for { raw, err := ln.Accept() if err != nil { if ctx.Err() != nil || errors.Is(err, net.ErrClosed) { return nil } + if isTemporaryAcceptError(err) { + tempDelay = nextAcceptRetryDelay(tempDelay) + s.log.Debug("Temporary accept error; retrying", zap.Duration("backoff", tempDelay), zap.Error(err)) + if !waitAcceptRetry(ctx, tempDelay) { + return nil + } + continue + } return fmt.Errorf("accept: %w", err) } + tempDelay = 0 wg.Add(1) go func() { @@ -428,7 +548,7 @@ func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, err if obfuscated { ln = transport.ObfuscatedListener(ln) } - return newCompatTransportListener(s.codec, ln).Accept() + return newCompatTransportListener(s.codec, ln, s.frameBudget).Accept() } // serveConn 处理单个传输连接：读帧并按 auth_key_id 分流。 @@ -444,6 +564,13 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) var current *Conn defer func() { + // A successful Recv transfers the frame reservation to serveConn. Release it only after + // this stack has stopped using b/plain; transport.Close may have raced us earlier and must + // not return that memory budget prematurely. + releaseInboundFrameOwnership(conn) + // 先同步关闭物理 socket，解除可能阻塞在 writer.Send 的 outbound actor； + // 再停止 logical Conn，避免 Close 等 actor 时反过来等到 write deadline。 + _ = conn.Close() if current != nil { s.conns.Unregister(current) current.Close() @@ -468,7 +595,7 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) var replay *bin.Buffer for { if replay != nil { - b.ResetTo(replay.Copy()) + b.ResetTo(replay.Buf) replay = nil } else { // 建立 session 前（current==nil，握手 + 首个加密消息之前）用较短的 handshakeTimeout @@ -491,11 +618,29 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) } if authKeyID == emptyAuthKeyID { + releaseHandshake, admitted := s.admission.tryAcquireHandshake() + if !admitted { + if err := s.sendProtoError(ctx, conn, codec.CodeTransportFlood); err != nil { + return err + } + return nil + } next, err := s.handleExchange(ctx, conn, &b) + releaseHandshake() if err != nil { return err } replay = next + // Exchange has finished consuming the original transport frame. Drop its + // potentially near-16MiB backing immediately. A replay frame is the gotd + // encrypted-frame copy and keeps the existing frame reservation until it is + // dispatched; a completed handshake has no surviving frame and can release now. + trimOversizedInboundBuffer(&b) + if replay == nil { + releaseInboundFrameOwnership(conn) + } else { + retainInboundFrameBackings(conn, replay) + } continue } @@ -513,7 +658,9 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) if err := s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound); err != nil { return err } - continue + // -404 对 TDesktop 是 terminal key failure；继续保留 socket 只会允许 + // 同一客户端反复触发 AuthKeyStore 查询。回包一次后立即断开。 + return nil } fetchedKey = &d } @@ -522,6 +669,20 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) if err != nil { return err } + trimOversizedInboundBuffer(&b) + trimOversizedInboundBuffer(&plain) + retainInboundFrameBackings(conn, &b, &plain) + } +} + +// maxRetainedConnBuffer keeps normal upload/download frames allocation-free while preventing one +// exceptional near-16MiB transport frame from pinning that capacity for the lifetime of a long +// connection. RPC bodies that outlive dispatch already own a budgeted Copy. +const maxRetainedConnBuffer = 2 << 20 + +func trimOversizedInboundBuffer(b *bin.Buffer) { + if b != nil && cap(b.Buf) > maxRetainedConnBuffer { + b.Buf = nil } } diff --git a/internal/mtprotoedge/server_test.go b/internal/mtprotoedge/server_test.go index 4fa899ed..1c129e26 100644 --- a/internal/mtprotoedge/server_test.go +++ b/internal/mtprotoedge/server_test.go @@ -358,7 +358,7 @@ func TestSamePortWebSocketTransportRoundTrip(t *testing.T) { serverDone := make(chan error, 1) go func() { - l := newCompatTransportListener(nil, wsLn) + l := newCompatTransportListener(nil, wsLn, newInboundFrameBudget(defaultInboundFrameGlobalMaxBytes)) defer func() { _ = l.Close() }() conn, err := l.Accept() diff --git a/internal/mtprotoedge/session_manager.go b/internal/mtprotoedge/session_manager.go index bdcf4e0e..3d172b73 100644 --- a/internal/mtprotoedge/session_manager.go +++ b/internal/mtprotoedge/session_manager.go @@ -4,7 +4,9 @@ import ( "context" "errors" "fmt" + "sort" "sync" + "sync/atomic" "time" "go.uber.org/zap" @@ -36,7 +38,8 @@ const ( // （typing/presence，不写 durable log）经 PushToUserTransient* 在未就绪时直接跳过、不入队， // 因此本队列被老化/溢出/重试耗尽丢弃时，丢的一定是 durable 条目——getDifference 以 // user_update_events 兜底补齐，丢弃不丢数据。 - pendingPushMaxAge = 60 * time.Second + pendingPushMaxAge = 60 * time.Second + defaultPendingPushMaxBytes = int64(256 << 20) // maxSessionsPerAuthKey：单个 raw auth_key 允许同时在线的 session 上限。telesrv 单 DC， // 一个客户端的全部连接（主连接 + 并发下载/上传）共享同一 auth_key、各用独立 session_id， // 故此上限须高于真实客户端单设备的并发连接峰值，否则会误杀活跃下载/主连接： @@ -52,10 +55,51 @@ const ( maxChannelIndexPerSession = 8192 ) +// forceCloseBatchTimeout is one deadline for a whole revoke/replace/eviction batch. Conn.Close +// already bounds its inbound-RPC wait, but calling ForceClose serially would multiply that bound +// by the number of sessions. The batch helper starts every close concurrently and waits at most +// this one shared interval. +const forceCloseBatchTimeout = rpcCloseWaitTimeout + +// maxForceCloseParallelism caps control-plane close goroutines even if a corrupted/runtime index +// hands a revoke path far more sessions than maxSessionsPerAuthKey. Every Conn's producer/RPC gate +// is closed synchronously before these workers start, so a stuck transport.Close cannot admit more +// memory while the bounded workers continue draining physical sockets in the background. +const maxForceCloseParallelism = 64 + type queuedPush struct { - t proto.MessageType - msg bin.Encoder - at time.Time + t proto.MessageType + encoded *encodedOutboundMessage + reservation *pendingPushReservation + at time.Time +} + +type pendingPushReservation struct { + budget *outboundTrackedBudget + bytes int + refs atomic.Int32 +} + +func (r *pendingPushReservation) retain() { + if r == nil { + return + } + if refs := r.refs.Add(1); refs <= 1 { + panic("mtprotoedge: retained released pending push reservation") + } +} + +func (r *pendingPushReservation) release() { + if r == nil { + return + } + refs := r.refs.Add(-1) + if refs < 0 { + panic("mtprotoedge: pending push reservation released more than retained") + } + if refs == 0 { + r.budget.release(r.bytes) + } } type sessionKey struct { @@ -85,6 +129,7 @@ type SessionManager struct { bySessionMembers map[sessionKey]map[int64]struct{} pending map[sessionKey][]queuedPush // updates-ready 前暂存的主动推送 flushing map[sessionKey]bool // 置位时暂存正在排空的 session；排空完成前推送继续进 pending 保序 + pendingBudget *outboundTrackedBudget // 未就绪 session 暂存 encoded body 的进程级上限 lifecycle SessionLifecycleObserver log *zap.Logger @@ -107,6 +152,7 @@ func NewSessionManager(log *zap.Logger) *SessionManager { bySessionMembers: make(map[sessionKey]map[int64]struct{}), pending: make(map[sessionKey][]queuedPush), flushing: make(map[sessionKey]bool), + pendingBudget: newOutboundTrackedBudget(defaultPendingPushMaxBytes), log: log, } } @@ -161,11 +207,13 @@ func (m *SessionManager) Register(c *Conn) { ) m.mu.Unlock() - if replaced != nil { - replaced.Close() - } - if evicted != nil { - evicted.Close() + // 同 identity 的新物理连接已经原子接管索引；立即关闭旧 transport，不能只停 + // actor 后让旧 FD/read goroutine 滞留到 read timeout。replacement 与 cap eviction + // 共用一个并发关闭批次，不能把每条 Conn 的 RPC 等待上界串行相加。 + if replaced != nil || evicted != nil { + if !forceCloseConnBatch([]*Conn{replaced, evicted}, forceCloseBatchTimeout) { + m.log.Warn("Session replacement/eviction close exceeded shared deadline") + } } } @@ -217,7 +265,12 @@ func (m *SessionManager) DestroySession(sessionID int64) bool { zap.Int("online", len(m.bySession)), ) m.mu.Unlock() - c.Close() + if !forceCloseConnBatch([]*Conn{c}, forceCloseBatchTimeout) { + m.log.Warn("Destroyed session close exceeded shared deadline", + zap.String("auth_key_id", sessionKeyLog(key.authKeyID)), + zap.Int64("session_id", sessionID), + ) + } if observer != nil && offlineUser != 0 { observer.SessionOffline(key.authKeyID, sessionID, offlineUser, lastForUser) } @@ -230,7 +283,7 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i key := sessionKey{authKeyID: authKeyID, sessionID: sessionID} c, ok := m.bySession[key] if !ok { - delete(m.pending, key) + m.deletePendingLocked(key) m.mu.Unlock() return false } @@ -243,7 +296,12 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i zap.Int("online", len(m.bySession)), ) m.mu.Unlock() - c.Close() + if !forceCloseConnBatch([]*Conn{c}, forceCloseBatchTimeout) { + m.log.Warn("Destroyed session close exceeded shared deadline", + zap.String("auth_key_id", sessionKeyLog(authKeyID)), + zap.Int64("session_id", sessionID), + ) + } if observer != nil && offlineUser != 0 { observer.SessionOffline(authKeyID, sessionID, offlineUser, lastForUser) } @@ -287,7 +345,7 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) { c.membershipsSynced.Store(false) // 身份变化即丢弃暂存推送：它们属于前一个账号，flush 给新账号是跨账号泄露。 // 同时取消进行中的排空（runFlush 还另有 owner 校验做批内兜底）。 - delete(m.pending, key) + m.deletePendingLocked(key) delete(m.flushing, key) } } @@ -298,7 +356,7 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) { m.clearChannelInterestsLocked(key) m.clearChannelMembershipsLocked(c, key) c.membershipsSynced.Store(false) - delete(m.pending, key) + m.deletePendingLocked(key) delete(m.flushing, key) } } @@ -399,7 +457,7 @@ func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8 m.clearChannelInterestsLocked(key) m.clearChannelMembershipsLocked(c, key) c.membershipsSynced.Store(false) - delete(m.pending, key) + m.deletePendingLocked(key) delete(m.flushing, key) c.userID.Store(0) c.userIDResolved.Store(false) @@ -459,8 +517,11 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int ) } m.mu.Unlock() - for _, c := range conns { - c.ForceClose() + if !forceCloseConnBatch(conns, forceCloseBatchTimeout) { + m.log.Warn("Revoked auth-key session close exceeded shared deadline", + zap.String("auth_key_id", sessionKeyLog(authKeyID)), + zap.Int("sessions", len(conns)), + ) } if observer != nil { for _, e := range events { @@ -493,8 +554,11 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc } observer := m.lifecycle m.mu.Unlock() - for _, c := range conns { - c.ForceClose() + if !forceCloseConnBatch(conns, forceCloseBatchTimeout) { + m.log.Warn("Raw auth-key session close exceeded shared deadline", + zap.String("auth_key_id", sessionKeyLog(authKeyID)), + zap.Int("sessions", len(conns)), + ) } if observer != nil { for _, e := range events { @@ -504,6 +568,99 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc return len(conns) } +// forceCloseConnBatch closes every producer/RPC gate first, then closes physical transports with a +// bounded worker set. Physical close and actor/RPC convergence share one batch deadline; the wait is +// never multiplied by the number of sessions. Workers may finish physical closes after the caller's +// deadline, but no timed-out Conn can enqueue more work in that interval. Nil/duplicate entries are +// removed so Register's replacement/eviction slots cannot close the same Conn twice. +func forceCloseConnBatch(conns []*Conn, timeout time.Duration) bool { + if len(conns) == 0 { + return true + } + unique := make([]*Conn, 0, len(conns)) + seen := make(map[*Conn]struct{}, len(conns)) + for _, c := range conns { + if c == nil { + continue + } + if _, ok := seen[c]; ok { + continue + } + seen[c] = struct{}{} + unique = append(unique, c) + } + if len(unique) == 0 { + return true + } + + // This phase is non-blocking and must precede transport.Close: it is the safety boundary if + // an implementation of transport.Conn.Close itself blocks past the batch deadline. + for _, c := range unique { + c.beginTerminalShutdown() + } + + workers := min(len(unique), maxForceCloseParallelism) + jobs := make(chan *Conn, len(unique)) + for _, c := range unique { + jobs <- c + } + close(jobs) + var closeWG sync.WaitGroup + closeWG.Add(workers) + for range workers { + go func() { + defer closeWG.Done() + for c := range jobs { + c.closeTransport() + } + }() + } + physicalDone := make(chan struct{}) + go func() { + closeWG.Wait() + close(physicalDone) + }() + + if timeout <= 0 { + return false + } + deadline := time.Now().Add(timeout) + timer := time.NewTimer(time.Until(deadline)) + defer timer.Stop() + select { + case <-physicalDone: + case <-timer.C: + return false + } + + // All physical close calls returned. Wait for memory-owning actor/RPC work using the same + // deadline; the first genuinely stuck Conn consumes the remaining allowance, not a fresh 5s. + for _, c := range unique { + remaining := time.Until(deadline) + if remaining <= 0 { + return false + } + if c.rpcScheduler != nil && !c.waitInboundShutdown(remaining) { + return false + } + if c.outboundDone == nil { + continue + } + remaining = time.Until(deadline) + if remaining <= 0 { + return false + } + wait := time.NewTimer(remaining) + select { + case <-c.outboundDone: + wait.Stop() + case <-wait.C: + return false + } + } + return true +} + // UnbindAuthKey 清理某业务 auth_key 下所有活跃连接的登录用户缓存。 func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int { m.mu.Lock() @@ -520,7 +677,7 @@ func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int { m.clearChannelMembershipsLocked(c, key) c.membershipsSynced.Store(false) // 授权解除后暂存推送属于已登出的账号，不能等下一个登录者置位时 flush 出去。 - delete(m.pending, key) + m.deletePendingLocked(key) delete(m.flushing, key) c.userIDResolved.Store(true) count++ @@ -595,7 +752,7 @@ func (m *SessionManager) runFlush(c *Conn, key sessionKey, owner int64, attempt } if c.userID.Load() != owner { // 排空期间发生登出/换号：剩余暂存属于旧账号，丢弃且不得发给新账号。 - delete(m.pending, key) + m.deletePendingLocked(key) delete(m.flushing, key) m.mu.Unlock() return @@ -613,38 +770,47 @@ func (m *SessionManager) runFlush(c *Conn, key sessionKey, owner int64, attempt // 每条发送前复查身份：登出/换号后 batch 的剩余条目不能继续发到已易主的连接。 if c.userID.Load() != owner { m.mu.Lock() - delete(m.pending, key) + m.deletePendingLocked(key) delete(m.flushing, key) m.mu.Unlock() + releaseQueuedPushes(batch[i:]) return } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - err := c.Send(ctx, item.t, item.msg) + // Pending entries are durable account updates. Shared body-budget pressure is not + // evidence that this socket is corrupt, so use the non-terminal enqueue path; after + // bounded retries, getDifference is the authoritative recovery path. + err := c.SendBestEffortEncoded(ctx, item.t, item.encoded, 5*time.Second) cancel() if err == nil { + item.release() continue } m.mu.Lock() if cur, ok := m.bySession[key]; !ok || cur != c || !m.flushing[key] || c.userID.Load() != owner { // 连接换代/取消/易主：剩余 batch 不属于当前连接当前账号，丢弃。 if c.userID.Load() != owner { - delete(m.pending, key) + m.deletePendingLocked(key) delete(m.flushing, key) } m.mu.Unlock() + releaseQueuedPushes(batch[i:]) return } rest := append(append([]queuedPush(nil), batch[i:]...), m.pending[key]...) if len(rest) > maxPendingPushesPerSession { // 与 queueLocked 溢出策略一致：丢最旧留最新，让 pts 空洞集中在最前端， // flush 首条即触发客户端 gap 检测，恢复路径最短。 - rest = rest[len(rest)-maxPendingPushesPerSession:] + dropped := len(rest) - maxPendingPushesPerSession + releaseQueuedPushes(rest[:dropped]) + rest = rest[dropped:] } m.pending[key] = rest if attempt+1 >= maxFlushAttempts { // 重试用尽：置位激活避免 idle 客户端永久断流；剩余暂存中的 durable 更新 // 由客户端后续 pts 空洞触发 getDifference 补齐。 c.receivesUpdates.Store(true) + m.deletePendingLocked(key) delete(m.flushing, key) m.mu.Unlock() m.log.Debug("Flush gave up after retries; activated with getDifference fallback", @@ -717,41 +883,61 @@ func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, session // PushToSession 向指定 session 推送一条消息。 func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error { - m.mu.Lock() + m.mu.RLock() c, key, ok, ambiguous := m.uniqueSessionLocked(sessionID) if ambiguous { - m.mu.Unlock() + m.mu.RUnlock() return ErrSessionAmbiguous } if !ok { - m.mu.Unlock() + m.mu.RUnlock() return ErrSessionNotFound } - if !c.receivesUpdates.Load() { - m.queueLocked(key, t, msg) - m.mu.Unlock() - return nil + ready := c.receivesUpdates.Load() + m.mu.RUnlock() + if ready { + return c.Send(ctx, t, msg) } - m.mu.Unlock() - return c.Send(ctx, t, msg) + return m.queueOrSendPrepared(ctx, key, t, msg) } // PushToSessionForAuthKey 向指定 raw auth_key_id + session_id 推送一条消息。 func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error { - m.mu.Lock() + m.mu.RLock() key := sessionKey{authKeyID: authKeyID, sessionID: sessionID} c, ok := m.bySession[key] + if !ok { + m.mu.RUnlock() + return ErrSessionNotFound + } + ready := c.receivesUpdates.Load() + m.mu.RUnlock() + if ready { + return c.Send(ctx, t, msg) + } + return m.queueOrSendPrepared(ctx, key, t, msg) +} + +func (m *SessionManager) queueOrSendPrepared(ctx context.Context, key sessionKey, t proto.MessageType, msg bin.Encoder) error { + encoded, reservation, err := m.preparePendingPush(ctx, msg) + if err != nil { + return err + } + defer reservation.release() + + m.mu.Lock() + c, ok := m.bySession[key] if !ok { m.mu.Unlock() return ErrSessionNotFound } if !c.receivesUpdates.Load() { - m.queueLocked(key, t, msg) + _ = m.queuePreparedLocked(key, t, encoded, reservation) m.mu.Unlock() return nil } m.mu.Unlock() - return c.Send(ctx, t, msg) + return c.SendEncoded(ctx, t, encoded) } // PushToSessionForAuthKeyImmediate 向指定 raw auth_key_id + session_id 立即推送一条消息。 @@ -782,7 +968,7 @@ func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, ex return m.pushToUser(ctx, userID, nil, excludeSessionID, t, msg) } -// PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定业务 auth_key + session。 +// PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定 raw auth_key + session。 func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) { return m.pushToUser(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg) } @@ -793,23 +979,28 @@ func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, use // 漏 temp-key 设备）。未就绪连接跳过、不进 pending——密聊消息 durable 在 qts 队列， // 离线设备靠 getDifference 补回（在线推送只是加速器）。c.userID 复查防跨账号泄露。 func (m *SessionManager) PushToUserAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder) (int, error) { - getEncoded := onceEncodedOutbound(msg) - return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, false, func(c *Conn) error { - if c.outbound == nil || c.outboundControl == nil { - return ErrConnClosed - } - encoded, err := getEncoded() - if err != nil { - return err - } - return c.SendEncoded(ctx, t, encoded) - }) + // Secret-chat qts is the durable source of truth, so online delivery is an accelerator just + // like account pts fan-out. Do not synchronously wait for every PFS/raw connection's socket. + return m.pushToBusinessAuthKeyBestEffort(ctx, userID, businessAuthKeyID, t, msg, 2*time.Second) } // PushToUserAuthKeyTransient 是 PushToUserAuthKey 的 transient（typing）best-effort 版本。 func (m *SessionManager) PushToUserAuthKeyTransient(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) { - getEncoded := onceEncodedOutbound(msg) - return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, true, func(c *Conn) error { + return m.pushToBusinessAuthKeyBestEffort(ctx, userID, businessAuthKeyID, t, msg, timeout) +} + +func (m *SessionManager) pushToBusinessAuthKeyBestEffort(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) { + getEncoded := onceEncodedOutbound(ctx, msg) + var deadline time.Time + if timeout > 0 { + deadline = time.Now().Add(timeout) + } + if ctx != nil { + if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) { + deadline = ctxDeadline + } + } + return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, func(c *Conn) error { if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed } @@ -817,11 +1008,18 @@ func (m *SessionManager) PushToUserAuthKeyTransient(ctx context.Context, userID if err != nil { return err } - return c.SendBestEffortEncoded(ctx, t, encoded, timeout) + remaining := timeout + if !deadline.IsZero() { + remaining = time.Until(deadline) + if remaining < 0 { + remaining = 0 + } + } + return c.SendBestEffortEncoded(ctx, t, encoded, remaining) }) } -func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, transient bool, send func(*Conn) error) (int, error) { +func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, send func(*Conn) error) (int, error) { m.mu.Lock() candidates := m.businessAuthKeyCandidatesLocked(businessAuthKeyID) conns := make([]*Conn, 0, len(candidates)) @@ -836,7 +1034,6 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64 conns = append(conns, c) } m.mu.Unlock() - _ = transient var firstErr error sent := 0 for _, c := range conns { @@ -845,6 +1042,18 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64 continue } if err := send(c); err != nil { + if errors.Is(err, ErrOutboundTrackedBudget) { + // Shared process pressure is not evidence that this particular socket is + // slow. Skip this online accelerator; durable qts/difference is the truth. + continue + } + if errors.Is(err, ErrOutboundQueueFull) { + c.dropSlowConsumer() + continue + } + if errors.Is(err, ErrConnClosed) { + continue + } if firstErr == nil { firstErr = err } @@ -856,7 +1065,7 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64 } func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) { - getEncoded := onceEncodedOutbound(msg) + getEncoded := onceEncodedOutbound(ctx, msg) return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, func(c *Conn) error { if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed @@ -875,7 +1084,7 @@ func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAu // 下一次状态变化重建，囤积过期 transient 既无意义又会被 pending 的老化/溢出/重试耗尽误当 // 「durable 兜底」丢弃。走 best-effort 发送，不阻塞调用方。 func (m *SessionManager) PushToUserTransientExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) { - getEncoded := onceEncodedOutbound(msg) + getEncoded := onceEncodedOutbound(ctx, msg) return m.pushToUserWithSender(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, false, func(c *Conn) error { if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed @@ -897,7 +1106,19 @@ func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Co } func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) { - getEncoded := onceEncodedOutbound(msg) + getEncoded := onceEncodedOutbound(ctx, msg) + // timeout 是整次 fan-out 的等待预算，不是每个 session 各自一份。健康连接始终先走 + // SendBestEffortEncoded 的非阻塞快路径；预算耗尽后 remaining=0，仍会尝试快路径， + // 但不会再为后续慢连接串行等待。 + var deadline time.Time + if timeout > 0 { + deadline = time.Now().Add(timeout) + } + if ctx != nil { + if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) { + deadline = ctxDeadline + } + } return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, func(c *Conn) error { if c.outbound == nil || c.outboundControl == nil { return ErrConnClosed @@ -906,18 +1127,25 @@ func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64, if err != nil { return err } - return c.SendBestEffortEncoded(ctx, t, encoded, timeout) + remaining := timeout + if !deadline.IsZero() { + remaining = time.Until(deadline) + if remaining < 0 { + remaining = 0 + } + } + return c.SendBestEffortEncoded(ctx, t, encoded, remaining) }) } -func onceEncodedOutbound(msg bin.Encoder) func() (*encodedOutboundMessage, error) { +func onceEncodedOutbound(ctx context.Context, msg bin.Encoder) func() (*encodedOutboundMessage, error) { var ( encoded *encodedOutboundMessage err error ) return func() (*encodedOutboundMessage, error) { if encoded == nil && err == nil { - encoded, err = encodeOutboundMessage(msg) + encoded, err = encodeOutboundMessageContext(ctx, msg) } return encoded, err } @@ -957,6 +1185,10 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, } m.mu.RUnlock() if needQueue { + // TL encoding and the process-wide pending-byte reservation may be expensive or + // briefly block on the global encode gate. Do both before taking SessionManager.mu, + // then share the immutable body across every not-ready session found by the re-scan. + pendingEncoded, pendingReservation, pendingErr := m.preparePendingPush(ctx, msg) // 写锁下完整重扫（读锁释放到此之间状态可能变化，以重扫结果为准）。 conns = conns[:0] queued, dropped, excluded, skipped = 0, 0, 0, 0 @@ -972,7 +1204,7 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, skipped++ continue } - if m.queueLocked(key, t, msg) { + if pendingErr == nil && m.queuePreparedLocked(key, t, pendingEncoded, pendingReservation) { queued++ if debug { m.log.Debug("Push queued (session not updates-ready)", @@ -996,6 +1228,15 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, conns = append(conns, c) } m.mu.Unlock() + if pendingReservation != nil { + pendingReservation.release() // drop producer ref; queued entries own the body now. + } + if pendingErr != nil && debug { + m.log.Debug("Drop pending pushes outside byte budget", + zap.Int64("user_id", userID), + zap.Error(pendingErr), + ) + } } var firstErr error @@ -1008,6 +1249,29 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, continue } if err := send(c); err != nil { + if errors.Is(err, ErrOutboundTrackedBudget) { + // Do not turn pressure owned by other sockets into a reconnect storm on + // healthy recipients. The durable event remains recoverable by difference. + dropped++ + continue + } + // 对 durable/best-effort fan-out，队列满意味着该 socket 已成为慢消费者。 + // 立即摘除并把它视为离线：不能让其错误把已经投递给健康 session 的 outbox + // 行整体重试。该 session 的 durable gap 由 getDifference 恢复。 + if errors.Is(err, ErrOutboundQueueFull) { + c.dropSlowConsumer() + if debug { + m.log.Debug("Drop slow outbound consumer", + zap.Int64("user_id", userID), + zap.String("auth_key_id", sessionKeyLog(c.authKeyID)), + zap.Int64("session_id", c.sessionID), + ) + } + continue + } + if errors.Is(err, ErrConnClosed) { + continue + } if firstErr == nil { firstErr = err } @@ -1055,6 +1319,24 @@ func (m *SessionManager) Online() int { return len(m.bySession) } +// ActiveRawAuthKeyIDs 返回当前物理连接实际使用的 raw auth_key_id 去重快照。 +// maintenance 用它保护“已建 key 但尚未登录”的长连接不被 orphan GC 删除；不能用 +// business/temp→perm key 替代，否则活跃 temp 连接仍可能误删。 +func (m *SessionManager) ActiveRawAuthKeyIDs() [][8]byte { + m.mu.RLock() + defer m.mu.RUnlock() + seen := make(map[[8]byte]struct{}, len(m.bySession)) + out := make([][8]byte, 0, len(m.byAuthKey)) + for key := range m.bySession { + if _, ok := seen[key.authKeyID]; ok { + continue + } + seen[key.authKeyID] = struct{}{} + out = append(out, key.authKeyID) + } + return out +} + // IsUserOnline returns whether userID has at least one active connection. func (m *SessionManager) IsUserOnline(userID int64) bool { if userID == 0 { @@ -1268,6 +1550,54 @@ func (m *SessionManager) OnlineChannelMemberUserIDsExcluding(channelID int64, ex return out } +// OnlineChannelIDsSnapshot returns every channel with at least one live joined-member session in +// strictly ascending order. The global SessionManager lock is held only while copying map keys; +// sorting and all recovery database work happen after unlock. The fixed saturation-recovery actor +// is the sole caller, so its exceptional-path temporary memory is one int64 slice (peak about 8*C +// bytes) rather than repeated O(C) scans under the connection/membership lock. +func (m *SessionManager) OnlineChannelIDsSnapshot() []int64 { + m.mu.RLock() + out := make([]int64, 0, len(m.byMemberChannel)) + for channelID, sessions := range m.byMemberChannel { + if channelID <= 0 || len(sessions) == 0 { + continue + } + live := false + for key := range sessions { + if _, ok := m.bySession[key]; ok { + live = true + break + } + } + if !live { + continue + } + out = append(out, channelID) + } + m.mu.RUnlock() + sort.Slice(out, func(i, j int) bool { return out[i] < out[j] }) + return out +} + +// OnlineChannelIDsAfter is retained for bounded diagnostics/tests. Production recovery takes one +// OnlineChannelIDsSnapshot per generation and slices it into pages, avoiding repeated full scans. +func (m *SessionManager) OnlineChannelIDsAfter(afterChannelID int64, limit int) []int64 { + if limit <= 0 { + return nil + } + const maxRecoveryPage = 4096 + if limit > maxRecoveryPage { + limit = maxRecoveryPage + } + all := m.OnlineChannelIDsSnapshot() + start := sort.Search(len(all), func(i int) bool { return all[i] > afterChannelID }) + end := start + limit + if end > len(all) { + end = len(all) + } + return all[start:end] +} + func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64, channelID int64, limit int) []int64 { if channelID == 0 { return nil @@ -1314,7 +1644,7 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 { m.clearChannelInterestsLocked(key) m.clearChannelMembershipsLocked(c, key) if dropPending { - delete(m.pending, key) + m.deletePendingLocked(key) } delete(m.flushing, key) return uid @@ -1430,8 +1760,11 @@ func (m *SessionManager) takePendingLocked(key sessionKey, ready bool) []queuedP now := time.Now() pending := make([]queuedPush, 0, len(q)) dropped := 0 - for _, item := range q { + for i := range q { + item := q[i] + q[i] = queuedPush{} if now.Sub(item.at) > pendingPushMaxAge { + item.release() dropped++ continue } @@ -1447,9 +1780,43 @@ func (m *SessionManager) takePendingLocked(key sessionKey, ready bool) []queuedP return pending } -// queueLocked 暂存一条主动推送，返回是否实际入队——stale 丢批分支会连同当前 -// 这条一起丢弃，调用方据此区分 queued/dropped 计数，避免投递日志失真。 -func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool { +// preparePendingPush encodes outside SessionManager.mu and reserves the one physical body before +// releasing the process-wide encode slot. Multiple not-ready sessions may then share this +// immutable body via reservation refs instead of encoding/copying it once per session. +func (m *SessionManager) preparePendingPush(ctx context.Context, msg bin.Encoder) (*encodedOutboundMessage, *pendingPushReservation, error) { + var ( + encoded *encodedOutboundMessage + bytes int + ) + err := withOutboundEncodeSlot(ctx, nil, func() error { + var err error + encoded, err = encodeOutboundMessageWithoutSlot(msg) + if err != nil { + return err + } + if encoded == nil { + return errors.New("nil encoded pending push") + } + bytes = len(encoded.body) + if bytes > maxOutboundBodyBytes { + return fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, bytes, maxOutboundBodyBytes) + } + if !m.pendingBudget.reserve(bytes) { + return ErrOutboundTrackedBudget + } + return nil + }) + if err != nil { + return nil, nil, err + } + reservation := &pendingPushReservation{budget: m.pendingBudget, bytes: bytes} + reservation.refs.Store(1) // producer ownership; queue entries retain below. + return encoded, reservation, nil +} + +// queuePreparedLocked 暂存一条已编码的主动推送，返回是否实际入队。 +// 调用方必须在锁外保持 reservation 的 producer ref，并在全部入队完成后 release。 +func (m *SessionManager) queuePreparedLocked(key sessionKey, t proto.MessageType, encoded *encodedOutboundMessage, reservation *pendingPushReservation) bool { q := m.pending[key] // 过期保护：最早一条暂存已超过 pendingPushMaxAge（session 迟迟未 ready）时，丢整批并 // 不再囤这条，记 trace。避免「登录后从不 getState」的连接长期占用 pending 内存。 @@ -1459,11 +1826,21 @@ func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bi zap.Int64("session_id", key.sessionID), zap.Int("dropped", len(q)), ) - delete(m.pending, key) + m.deletePendingLocked(key) return false } - push := queuedPush{t: t, msg: msg, at: time.Now()} + if encoded == nil || reservation == nil { + return false + } + reservation.retain() + push := queuedPush{ + t: t, + encoded: encoded, + reservation: reservation, + at: time.Now(), + } if len(q) >= maxPendingPushesPerSession { + q[0].release() copy(q, q[1:]) q[len(q)-1] = push m.pending[key] = q @@ -1473,6 +1850,22 @@ func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bi return true } +// queueLocked remains as a test/internal single-target convenience. Production fan-out prepares +// outside m.mu and calls queuePreparedLocked so TL encoding never serializes the session registry. +func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool { + encoded, reservation, err := m.preparePendingPush(context.Background(), msg) + if err != nil { + m.log.Debug("Drop pending push outside byte budget", + zap.String("auth_key_id", sessionKeyLog(key.authKeyID)), + zap.Int64("session_id", key.sessionID), + zap.Error(err), + ) + return false + } + defer reservation.release() + return m.queuePreparedLocked(key, t, encoded, reservation) +} + func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey, bool, bool) { set := m.bySessionID[sessionID] if len(set) == 0 { @@ -1490,7 +1883,7 @@ func (m *SessionManager) uniqueSessionLocked(sessionID int64) (*Conn, sessionKey func (m *SessionManager) dropPendingBySessionLocked(sessionID int64) { for key := range m.pending { if key.sessionID == sessionID { - delete(m.pending, key) + m.deletePendingLocked(key) } } } @@ -1528,7 +1921,7 @@ func (m *SessionManager) sweepStalePending() { if len(q) == 0 || now.Sub(q[0].at) <= pendingPushMaxAge { continue } - delete(m.pending, key) + m.deletePendingLocked(key) dropped++ } if dropped > 0 { @@ -1536,6 +1929,27 @@ func (m *SessionManager) sweepStalePending() { } } +func (q *queuedPush) release() { + if q == nil { + return + } + reservation := q.reservation + *q = queuedPush{} + reservation.release() +} + +func releaseQueuedPushes(q []queuedPush) { + for i := range q { + q[i].release() + } +} + +func (m *SessionManager) deletePendingLocked(key sessionKey) { + q := m.pending[key] + delete(m.pending, key) + releaseQueuedPushes(q) +} + func addConnIndex[K comparable](idx map[K]map[int64]*Conn, key K, sessionID int64, c *Conn) { set := idx[key] if set == nil { @@ -1630,7 +2044,7 @@ func shouldExcludeSession(c *Conn, excludeAuthKeyID *[8]byte, excludeSessionID i if excludeAuthKeyID == nil || *excludeAuthKeyID == ([8]byte{}) { return true } - return connUsesBusinessAuthKey(c, *excludeAuthKeyID) + return c.authKeyID == *excludeAuthKeyID } func sessionKeyLog(id [8]byte) string { diff --git a/internal/mtprotoedge/session_manager_test.go b/internal/mtprotoedge/session_manager_test.go index 6153355b..3d3d5bbc 100644 --- a/internal/mtprotoedge/session_manager_test.go +++ b/internal/mtprotoedge/session_manager_test.go @@ -3,6 +3,8 @@ package mtprotoedge import ( "context" "errors" + "sync" + "sync/atomic" "testing" "time" @@ -27,6 +29,35 @@ type closeCountingTransport struct { closes int } +type slowCloseTransport struct { + delay time.Duration + release <-chan struct{} + done chan struct{} + once sync.Once + closes atomic.Int32 +} + +func newSlowCloseTransport(delay time.Duration, release <-chan struct{}) *slowCloseTransport { + return &slowCloseTransport{delay: delay, release: release, done: make(chan struct{})} +} + +func (*slowCloseTransport) Send(context.Context, *bin.Buffer) error { + return errors.New("test transport send") +} +func (*slowCloseTransport) Recv(context.Context, *bin.Buffer) error { + return errors.New("test transport recv") +} +func (t *slowCloseTransport) Close() error { + t.closes.Add(1) + if t.release != nil { + <-t.release + } else if t.delay > 0 { + time.Sleep(t.delay) + } + t.once.Do(func() { close(t.done) }) + return nil +} + func (t *closeCountingTransport) Send(context.Context, *bin.Buffer) error { return errors.New("test transport send") } @@ -78,6 +109,43 @@ func TestSessionManagerRegistry(t *testing.T) { } } +func TestSessionManagerReplacementClosesOldPhysicalTransport(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + raw := [8]byte{1, 2, 3} + oldTransport := &closeCountingTransport{} + old := &Conn{sessionID: 42, authKeyID: raw, transport: oldTransport} + replacement := &Conn{sessionID: 42, authKeyID: raw} + + sm.Register(old) + sm.Register(replacement) + if oldTransport.closes != 1 { + t.Fatalf("old transport closes = %d, want 1", oldTransport.closes) + } + // 旧 serveConn 稍后退出时不得把 replacement 从索引删掉。 + sm.Unregister(old) + if got, ok := sm.bySession[sessionKey{authKeyID: raw, sessionID: 42}]; !ok || got != replacement { + t.Fatal("old unregister removed the replacement connection") + } +} + +func TestSessionManagerDestroyClosesPhysicalTransport(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + raw := [8]byte{4, 5, 6} + physical := &closeCountingTransport{} + c := &Conn{sessionID: 77, authKeyID: raw, transport: physical} + sm.Register(c) + + if !sm.DestroySessionForAuthKey(raw, 77) { + t.Fatal("DestroySessionForAuthKey returned false") + } + if physical.closes != 1 { + t.Fatalf("destroyed transport closes = %d, want 1", physical.closes) + } + if sm.Online() != 0 { + t.Fatalf("online after destroy = %d, want 0", sm.Online()) + } +} + func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) { sm := NewSessionManager(zaptest.NewLogger(t)) const userID = int64(100) @@ -88,6 +156,7 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) { outbound: make(chan outboundOp, 1), outboundControl: make(chan outboundOp, 1), outboundStop: make(chan struct{}), + metrics: NopMetrics{}, } c.userID.Store(userID) c.userIDResolved.Store(true) @@ -115,6 +184,117 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) { } } +func TestSessionManagerPendingFanoutSharesOneEncodedBodyAndBudget(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + const userID = int64(102) + keys := make([]sessionKey, 0, 2) + for i := 0; i < 2; i++ { + c := &Conn{sessionID: int64(i + 1), authKeyID: [8]byte{byte(i + 1)}} + c.userID.Store(userID) + c.userIDResolved.Store(true) + sm.Register(c) + keys = append(keys, connSessionKey(c)) + } + + encodes := 0 + msg := &countingOutboundEncoder{count: &encodes} + sent, err := sm.PushToUserExceptSession(context.Background(), userID, 0, proto.MessageFromServer, msg) + if err != nil { + t.Fatalf("push: %v", err) + } + if sent != 2 || encodes != 1 { + t.Fatalf("pending fanout = sent:%d encodes:%d, want 2/1", sent, encodes) + } + + sm.mu.Lock() + first := sm.pending[keys[0]][0] + second := sm.pending[keys[1]][0] + if first.encoded != second.encoded || first.reservation != second.reservation { + sm.mu.Unlock() + t.Fatal("pending sessions did not share encoded body/reservation") + } + wantBytes := int64(len(first.encoded.body)) + sm.deletePendingLocked(keys[0]) + if got := sm.pendingBudget.snapshot(); got != wantBytes { + sm.mu.Unlock() + t.Fatalf("budget after first session drop = %d, want shared body %d", got, wantBytes) + } + sm.deletePendingLocked(keys[1]) + sm.mu.Unlock() + if got := sm.pendingBudget.snapshot(); got != 0 { + t.Fatalf("budget after last session drop = %d, want 0", got) + } +} + +func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + const userID = int64(101) + + // 三个满队列模拟三个慢设备；没有 outbound actor，确保队列在测试期间不会自行排空。 + slow := make([]*Conn, 0, 3) + for i := 0; i < 3; i++ { + tr := &closeCountingTransport{} + c := &Conn{ + sessionID: int64(i + 1), + authKeyID: [8]byte{byte(i + 1)}, + transport: tr, + metrics: NopMetrics{}, + outbound: make(chan outboundOp, 1), + outboundControl: make(chan outboundOp, 1), + outboundStop: make(chan struct{}), + } + c.outbound <- outboundOp{} + c.userID.Store(userID) + c.userIDResolved.Store(true) + c.receivesUpdates.Store(true) + sm.Register(c) + slow = append(slow, c) + } + + healthy := &Conn{ + sessionID: 99, + authKeyID: [8]byte{99}, + metrics: NopMetrics{}, + outbound: make(chan outboundOp, 1), + outboundControl: make(chan outboundOp, 1), + outboundStop: make(chan struct{}), + } + healthy.userID.Store(userID) + healthy.userIDResolved.Store(true) + healthy.receivesUpdates.Store(true) + sm.Register(healthy) + + const budget = 40 * time.Millisecond + start := time.Now() + sent, err := sm.PushToUserExceptSessionBestEffort( + context.Background(), userID, 0, proto.MessageFromServer, &tg.UpdatesTooLong{}, budget, + ) + elapsed := time.Since(start) + if err != nil { + t.Fatalf("push: %v", err) + } + if sent != 1 { + t.Fatalf("sent = %d, want only healthy session", sent) + } + if elapsed >= 3*budget { + t.Fatalf("fan-out waited %v; want one shared %v budget, not one per slow session", elapsed, budget) + } + if got := len(healthy.outbound); got != 1 { + t.Fatalf("healthy queued ops = %d, want 1", got) + } + if healthy.terminal.Load() { + t.Fatal("healthy session was terminalized") + } + for i, c := range slow { + if !c.terminal.Load() { + t.Fatalf("slow session %d was not terminalized", i) + } + if tr := c.transport.(*closeCountingTransport); tr.closes != 1 { + t.Fatalf("slow session %d transport closes = %d, want 1", i, tr.closes) + } + } +} + func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) { sm := NewSessionManager(zaptest.NewLogger(t)) raw1 := [8]byte{1} @@ -130,6 +310,9 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) { } sm.BindAuthKeyForSession(raw1, 42, perm1) + // 两条 PFS/raw 连接可以解析到同一业务 perm key 且复用同一个 session_id； + // 精确排除必须只匹配 raw1，不能按 business key 把 raw2 一并排除。 + sm.BindAuthKeyForSession(raw2, 42, perm1) sm.BindUserForAuthKey(raw1, 42, 100) sm.BindUserForAuthKey(raw2, 42, 200) @@ -148,7 +331,7 @@ func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) { sm.BindUserForAuthKey(raw1, 42, 300) sm.BindUserForAuthKey(raw2, 42, 300) - sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, perm1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{}) + sent, err := sm.PushToUserExceptAuthKeySession(context.Background(), 300, raw1, 42, proto.MessageFromServer, &tg.UpdatesTooLong{}) if err != nil { t.Fatalf("push except scoped session: %v", err) } @@ -213,6 +396,151 @@ func TestSessionManagerCloseSessionsForBusinessAuthKeyClosesBoundTempAndRaw(t *t } } +func TestSessionManagerCloseSessionsRunsSlowPhysicalClosesConcurrently(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + business := [8]byte{9, 9, 9} + const sessions = 8 + const closeDelay = 75 * time.Millisecond + transports := make([]*slowCloseTransport, 0, sessions) + for i := 0; i < sessions; i++ { + raw := [8]byte{byte(i + 1)} + tr := newSlowCloseTransport(closeDelay, nil) + c := &Conn{sessionID: int64(i + 1), authKeyID: raw, transport: tr} + sm.Register(c) + sm.BindAuthKeyForSession(raw, c.sessionID, business) + transports = append(transports, tr) + } + + started := time.Now() + if got := sm.CloseSessionsForBusinessAuthKey(business); got != sessions { + t.Fatalf("closed sessions = %d, want %d", got, sessions) + } + elapsed := time.Since(started) + // A serial implementation takes ~600ms. Leave ample Windows/CI scheduling margin while + // still proving that the per-Conn delay is not multiplied by the session count. + if elapsed >= 4*closeDelay { + t.Fatalf("batch close elapsed = %v, want concurrent closes near %v", elapsed, closeDelay) + } + for i, tr := range transports { + select { + case <-tr.done: + default: + t.Fatalf("transport %d close had not completed when batch returned", i) + } + if got := tr.closes.Load(); got != 1 { + t.Fatalf("transport %d closes = %d, want 1", i, got) + } + } +} + +func TestSessionManagerCloseRawSessionsExceptRunsConcurrentlyAndPreservesExcluded(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + raw := [8]byte{6, 6, 6} + const sessions = 7 + const excludedSession = int64(4) + const closeDelay = 60 * time.Millisecond + transports := make([]*slowCloseTransport, 0, sessions) + for i := 0; i < sessions; i++ { + tr := newSlowCloseTransport(closeDelay, nil) + c := &Conn{sessionID: int64(i + 1), authKeyID: raw, transport: tr} + sm.Register(c) + transports = append(transports, tr) + } + + started := time.Now() + if got, want := sm.CloseSessionsForRawAuthKeyExcept(raw, excludedSession), sessions-1; got != want { + t.Fatalf("closed sessions = %d, want %d", got, want) + } + if elapsed := time.Since(started); elapsed >= 4*closeDelay { + t.Fatalf("raw-key batch close elapsed = %v, want concurrent closes near %v", elapsed, closeDelay) + } + for i, tr := range transports { + sessionID := int64(i + 1) + if sessionID == excludedSession { + if got := tr.closes.Load(); got != 0 { + t.Fatalf("excluded transport closes = %d, want 0", got) + } + continue + } + select { + case <-tr.done: + default: + t.Fatalf("transport for session %d had not closed", sessionID) + } + } + if _, ok := sm.bySession[sessionKey{authKeyID: raw, sessionID: excludedSession}]; !ok { + t.Fatal("excluded session was removed from the registry") + } + // Clean up the deliberately preserved connection without making the assertion path depend + // on test process teardown. + if !sm.DestroySessionForAuthKey(raw, excludedSession) { + t.Fatal("cleanup destroy of excluded session failed") + } +} + +func TestForceCloseBatchTimeoutStillClosesProducerAndRPCGates(t *testing.T) { + release := make(chan struct{}) + const sessions = 4 + scheduler := newInboundRPCScheduler(1, 16, 1<<20) + defer scheduler.stop(time.Second) + conns := make([]*Conn, 0, sessions) + transports := make([]*slowCloseTransport, 0, sessions) + for i := 0; i < sessions; i++ { + tr := newSlowCloseTransport(0, release) + c := &Conn{ + transport: tr, + metrics: NopMetrics{}, + outbound: make(chan outboundOp, 1), + outboundControl: make(chan outboundOp, 1), + outboundStop: make(chan struct{}), + } + c.startInboundRPCScheduler(scheduler, 1, 1, time.Second) + if err := c.enqueueInboundRPC(context.Background(), inboundRPC{ + method: "shutdown.budget", + size: 32, + }); err != nil { + t.Fatalf("enqueue queued RPC %d: %v", i, err) + } + conns = append(conns, c) + transports = append(transports, tr) + } + + started := time.Now() + if completed := forceCloseConnBatch(conns, 40*time.Millisecond); completed { + t.Fatal("blocked transport close batch unexpectedly completed") + } + if elapsed := time.Since(started); elapsed > 250*time.Millisecond { + t.Fatalf("timed batch close blocked for %v", elapsed) + } + for i, c := range conns { + if !c.terminal.Load() { + t.Fatalf("connection %d producer gate remains open after batch timeout", i) + } + select { + case <-c.outboundStop: + default: + t.Fatalf("connection %d outbound stop was not published", i) + } + select { + case <-c.rpcRootCtx.Done(): + default: + t.Fatalf("connection %d RPC root remains open after batch timeout", i) + } + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("RPC budget after batch gate close = tasks:%d bytes:%d, want zero", tasks, bytes) + } + + close(release) + for i, tr := range transports { + select { + case <-tr.done: + case <-time.After(time.Second): + t.Fatalf("transport %d did not finish after release", i) + } + } +} + func TestSessionManagerBusinessAuthKeyIndexTracksRebind(t *testing.T) { sm := NewSessionManager(zaptest.NewLogger(t)) raw := [8]byte{1} @@ -238,6 +566,63 @@ func TestSessionManagerBusinessAuthKeyIndexTracksRebind(t *testing.T) { } } +func TestPushToUserAuthKeyUsesOneDeadlineAndDropsOnlySlowPFSConnections(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + business := [8]byte{9, 9} + const userID = int64(100) + newConn := func(raw [8]byte, sessionID int64, queueFull bool) (*Conn, *closeCountingTransport) { + transport := &closeCountingTransport{} + c := &Conn{ + authKeyID: raw, + sessionID: sessionID, + metrics: NopMetrics{}, + transport: transport, + outbound: make(chan outboundOp, 1), + outboundControl: make(chan outboundOp, 1), + outboundStop: make(chan struct{}), + } + c.receivesUpdates.Store(true) + if queueFull { + c.outbound <- outboundOp{} + } + sm.Register(c) + sm.BindAuthKeyForSession(raw, sessionID, business) + sm.BindUserForAuthKey(raw, sessionID, userID) + return c, transport + } + + slowOne, slowOneTransport := newConn([8]byte{1}, 11, true) + slowTwo, slowTwoTransport := newConn([8]byte{2}, 12, true) + healthy, healthyTransport := newConn([8]byte{3}, 13, false) + + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + started := time.Now() + sent, err := sm.PushToUserAuthKey(ctx, userID, business, proto.MessageFromServer, &tg.UpdatesTooLong{}) + elapsed := time.Since(started) + if err != nil { + t.Fatalf("PushToUserAuthKey: %v", err) + } + if sent != 1 { + t.Fatalf("sent = %d, want only healthy connection", sent) + } + if elapsed > 100*time.Millisecond { + t.Fatalf("elapsed = %v, want one shared deadline rather than per-session waits", elapsed) + } + if !slowOne.terminal.Load() || !slowTwo.terminal.Load() || slowOneTransport.closes != 1 || slowTwoTransport.closes != 1 { + t.Fatalf("slow connections not terminal/closed: one=%v/%d two=%v/%d", + slowOne.terminal.Load(), slowOneTransport.closes, slowTwo.terminal.Load(), slowTwoTransport.closes) + } + if healthy.terminal.Load() || healthyTransport.closes != 0 { + t.Fatalf("healthy connection was dropped: terminal=%v closes=%d", healthy.terminal.Load(), healthyTransport.closes) + } + select { + case <-healthy.outbound: + default: + t.Fatal("healthy PFS connection did not receive best-effort enqueue") + } +} + func TestSessionManagerChannelInterestIndex(t *testing.T) { sm := NewSessionManager(zaptest.NewLogger(t)) raw := [8]byte{1, 2, 3} @@ -368,8 +753,16 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) { select { case op := <-c.outbound: - if op.msg != msg { - t.Fatalf("enqueued msg = %T, want original update", op.msg) + defer op.releaseReservation(c.outboundTrackedBudget) + if op.encoded == nil { + t.Fatal("immediate push did not retain its encoded body") + } + var got tg.UpdateShort + if err := got.Decode(&bin.Buffer{Buf: op.encoded.body}); err != nil { + t.Fatalf("decode enqueued update: %v", err) + } + if _, ok := got.Update.(*tg.UpdateLoginToken); !ok || got.Date != msg.Date { + t.Fatalf("enqueued update = %+v, want login token date %d", got, msg.Date) } case <-time.After(time.Second): t.Fatal("immediate push was not enqueued") @@ -383,6 +776,126 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) { } } +func TestPendingPushBodiesUseGlobalByteBudgetAndReleaseOnDrop(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000} + encoded, err := encodeOutboundMessage(msg) + if err != nil { + t.Fatalf("encode pending fixture: %v", err) + } + sm.pendingBudget = newOutboundTrackedBudget(int64(len(encoded.body))) + key := sessionKey{authKeyID: [8]byte{9}, sessionID: 77} + + sm.mu.Lock() + first := sm.queueLocked(key, proto.MessageFromServer, msg) + second := sm.queueLocked(key, proto.MessageFromServer, msg) + sm.mu.Unlock() + if !first || second { + t.Fatalf("pending queue results = first %v second %v, want true/false at byte cap", first, second) + } + if got := sm.pendingBudget.snapshot(); got != int64(len(encoded.body)) { + t.Fatalf("pending body budget = %d, want %d", got, len(encoded.body)) + } + + sm.mu.Lock() + sm.deletePendingLocked(key) + sm.mu.Unlock() + if got := sm.pendingBudget.snapshot(); got != 0 { + t.Fatalf("pending body budget after drop = %d, want zero", got) + } +} + +func TestPendingFlushGlobalBodyPressureDoesNotTerminateHealthyConnection(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + key := sessionKey{authKeyID: [8]byte{6}, sessionID: 66} + c := &Conn{ + authKeyID: key.authKeyID, + sessionID: key.sessionID, + outbound: make(chan outboundOp, 1), + outboundControl: make(chan outboundOp, 1), + outboundStop: make(chan struct{}), + metrics: NopMetrics{}, + outboundTrackedBudget: newOutboundTrackedBudget(1), + } + const userID = int64(606) + c.userID.Store(userID) + c.userIDResolved.Store(true) + sm.Register(c) + + msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000} + sm.mu.Lock() + if !sm.queueLocked(key, proto.MessageFromServer, msg) { + sm.mu.Unlock() + t.Fatal("queue pending push") + } + sm.flushing[key] = true + sm.mu.Unlock() + + // Enter at the final retry so the test exercises the durable-difference fallback without + // waiting for the production backoff timer. + sm.runFlush(c, key, userID, maxFlushAttempts-1) + if c.terminal.Load() { + t.Fatal("shared body pressure terminated a healthy pending-flush connection") + } + if !c.receivesUpdates.Load() { + t.Fatal("pending flush did not activate difference fallback after bounded retries") + } + if got := sm.pendingBudget.snapshot(); got != 0 { + t.Fatalf("pending budget after fallback = %d, want zero", got) + } +} + +func TestPendingPushBudgetSurvivesTakeAndReturnsAcrossOverflowAndUnregister(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000} + encoded, err := encodeOutboundMessage(msg) + if err != nil { + t.Fatalf("encode pending fixture: %v", err) + } + bytesPerPush := int64(len(encoded.body)) + sm.pendingBudget = newOutboundTrackedBudget(bytesPerPush * (maxPendingPushesPerSession + 8)) + key := sessionKey{authKeyID: [8]byte{7}, sessionID: 55} + c := &Conn{authKeyID: key.authKeyID, sessionID: key.sessionID} + sm.Register(c) + + sm.mu.Lock() + for i := 0; i < maxPendingPushesPerSession+5; i++ { + if !sm.queueLocked(key, proto.MessageFromServer, msg) { + sm.mu.Unlock() + t.Fatalf("queue pending push %d unexpectedly failed", i) + } + } + if got, want := sm.pendingBudget.snapshot(), bytesPerPush*maxPendingPushesPerSession; got != want { + sm.mu.Unlock() + t.Fatalf("budget after overflow replacement = %d, want %d", got, want) + } + batch := sm.takePendingLocked(key, true) + sm.mu.Unlock() + if len(batch) != maxPendingPushesPerSession { + t.Fatalf("taken pending pushes = %d, want %d", len(batch), maxPendingPushesPerSession) + } + // take transfers ownership to runFlush; deleting the map entry must not release bodies while + // the batch still references them. + if got, want := sm.pendingBudget.snapshot(), bytesPerPush*maxPendingPushesPerSession; got != want { + t.Fatalf("budget after take = %d, want transferred ownership %d", got, want) + } + releaseQueuedPushes(batch) + if got := sm.pendingBudget.snapshot(); got != 0 { + t.Fatalf("budget after taken batch release = %d, want 0", got) + } + + sm.mu.Lock() + if !sm.queueLocked(key, proto.MessageFromServer, msg) { + sm.mu.Unlock() + t.Fatal("queue before unregister failed") + } + sm.mu.Unlock() + sm.Unregister(c) + if got := sm.pendingBudget.snapshot(); got != 0 { + t.Fatalf("budget after unregister = %d, want 0", got) + } +} + // TestSessionManagerPush 验证主动推送端到端：两个 client 连接握手并建立 session 后， // server 经 PushToSession / PushToUser 主动向其推送，client 收到。 func TestSessionManagerPush(t *testing.T) { diff --git a/internal/mtprotoedge/session_membership_gen_test.go b/internal/mtprotoedge/session_membership_gen_test.go index 47811cb0..76e142b2 100644 --- a/internal/mtprotoedge/session_membership_gen_test.go +++ b/internal/mtprotoedge/session_membership_gen_test.go @@ -1,12 +1,62 @@ package mtprotoedge import ( + "slices" "testing" "time" "go.uber.org/zap/zaptest" ) +func TestOnlineChannelIDsSnapshotAndDiagnosticPagesStableAscending(t *testing.T) { + sm := NewSessionManager(zaptest.NewLogger(t)) + raw := [8]byte{4, 5, 6} + c := &Conn{sessionID: 77, authKeyID: raw} + sm.Register(c) + sm.BindUserForAuthKey(raw, 77, 100) + sm.SetSessionChannelMemberships(raw, 77, 100, []int64{50, 10, 30, 20, 40}, sm.ChannelMembershipGeneration(raw, 77)) + want := []int64{10, 20, 30, 40, 50} + snapshot := sm.OnlineChannelIDsSnapshot() + if !slices.Equal(snapshot, want) { + t.Fatalf("online channel snapshot = %v, want %v", snapshot, want) + } + + var got []int64 + after := int64(0) + for { + page := sm.OnlineChannelIDsAfter(after, 2) + if len(page) == 0 { + break + } + for _, channelID := range page { + if channelID <= after { + t.Fatalf("page %v not strictly after cursor %d", page, after) + } + after = channelID + got = append(got, channelID) + } + } + if !slices.Equal(got, want) { + t.Fatalf("paged online channels = %v, want %v", got, want) + } + // The recovery actor owns a stable copy: later membership changes are visible to the next + // generation, not spliced into the in-flight sorted snapshot. + sm.AddUserChannelMembership(100, 5) + if !slices.Equal(snapshot, want) { + t.Fatalf("owned snapshot mutated after membership insert: %v", snapshot) + } + if current := sm.OnlineChannelIDsSnapshot(); !slices.Equal(current, []int64{5, 10, 20, 30, 40, 50}) { + t.Fatalf("next online channel snapshot = %v", current) + } + + // Removing the only live session must immediately remove all channel ids from the recovery + // enumeration; stale membership map entries are never enough without a live bySession key. + sm.Unregister(c) + if got := sm.OnlineChannelIDsAfter(0, 10); len(got) != 0 { + t.Fatalf("online channels after unregister = %v, want empty", got) + } +} + // TestSetSessionChannelMembershipsDetectsConcurrentIncrementalUpdates 验证全量 // membership 同步的丢失更新防护：同步方在读持久成员列表前采样修订号，读取窗口内 // 若发生增量 join/leave（另一设备操作经 Add/RemoveUserChannelMembership 落索引）， @@ -67,13 +117,18 @@ func TestRegisterEvictsOldestSessionAtCap(t *testing.T) { base := time.Unix(1_700_000_000, 0) const oldestSession = int64(100) + oldestTransport := &closeCountingTransport{} for i := 0; i < maxSessionsPerAuthKey; i++ { sid := int64(i + 1) created := base.Add(time.Duration(i+1) * time.Second) if sid == oldestSession { created = base // 唯一早于所有其它连接的时间戳，且故意不在注册顺序首位。 } - sm.Register(&Conn{sessionID: sid, authKeyID: raw, createdAt: created}) + c := &Conn{sessionID: sid, authKeyID: raw, createdAt: created} + if sid == oldestSession { + c.transport = oldestTransport + } + sm.Register(c) } sm.Register(&Conn{sessionID: 9999, authKeyID: raw, createdAt: base.Add(time.Hour)}) @@ -92,4 +147,7 @@ func TestRegisterEvictsOldestSessionAtCap(t *testing.T) { if total != maxSessionsPerAuthKey { t.Fatalf("sessions for auth key = %d, want cap %d", total, maxSessionsPerAuthKey) } + if oldestTransport.closes != 1 { + t.Fatalf("evicted transport closes = %d, want 1", oldestTransport.closes) + } } diff --git a/internal/mtprotoedge/shutdown_gate_test.go b/internal/mtprotoedge/shutdown_gate_test.go new file mode 100644 index 00000000..f953ee3a --- /dev/null +++ b/internal/mtprotoedge/shutdown_gate_test.go @@ -0,0 +1,65 @@ +package mtprotoedge + +import ( + "testing" + "time" +) + +func TestTerminalFailurePathsCloseGatesBeforeBlockingTransportClose(t *testing.T) { + tests := []struct { + name string + run func(*Conn) + }{ + {name: "write failure", run: (*Conn).failTransport}, + {name: "slow consumer", run: (*Conn).dropSlowConsumer}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + release := make(chan struct{}) + tr := newSlowCloseTransport(0, release) + scheduler := newInboundRPCScheduler(1, 1, 1024) + defer scheduler.stop(time.Second) + c := &Conn{ + transport: tr, + metrics: NopMetrics{}, + outbound: make(chan outboundOp, 1), + outboundControl: make(chan outboundOp, 1), + outboundStop: make(chan struct{}), + } + c.startInboundRPCScheduler(scheduler, 1, 1, time.Second) + returned := make(chan struct{}) + go func() { + tt.run(c) + close(returned) + }() + + deadline := time.Now().Add(time.Second) + for tr.closes.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if tr.closes.Load() == 0 { + t.Fatal("terminal path did not enter transport.Close") + } + if !c.terminal.Load() { + t.Fatal("producer terminal gate was not published before blocking Close") + } + select { + case <-c.outboundStop: + default: + t.Fatal("outbound stop was not published before blocking Close") + } + select { + case <-c.rpcRootCtx.Done(): + default: + t.Fatal("RPC root was not canceled before blocking Close") + } + + close(release) + select { + case <-returned: + case <-time.After(time.Second): + t.Fatal("terminal path did not return after transport release") + } + }) + } +} diff --git a/internal/mtprotoedge/structural_limits_test.go b/internal/mtprotoedge/structural_limits_test.go new file mode 100644 index 00000000..10904369 --- /dev/null +++ b/internal/mtprotoedge/structural_limits_test.go @@ -0,0 +1,206 @@ +package mtprotoedge + +import ( + "context" + "strings" + "testing" + + "github.com/gotd/td/bin" + "github.com/gotd/td/mt" + "github.com/gotd/td/proto" + "go.uber.org/zap/zaptest" +) + +func TestContainerMessageCountAndServiceVectorCaps(t *testing.T) { + var container bin.Buffer + container.PutID(proto.MessageContainerTypeID) + container.PutInt(maxContainerMessages) + if got, err := containerMessageCount(&container); err != nil || got != maxContainerMessages { + t.Fatalf("container count = %d/%v, want %d/nil", got, err, maxContainerMessages) + } + container.Buf[4]++ + if got, err := containerMessageCount(&container); err != nil || got != maxContainerMessages+1 { + t.Fatalf("oversized container preflight = %d/%v, want %d/nil", got, err, maxContainerMessages+1) + } + + ack := mt.MsgsAck{MsgIDs: make([]int64, maxServiceMessageIDs)} + var encoded bin.Buffer + if err := ack.Encode(&encoded); err != nil { + t.Fatalf("encode msgs_ack: %v", err) + } + if err := validateFirstVectorCount(&encoded, maxServiceMessageIDs); err != nil { + t.Fatalf("service vector at cap: %v", err) + } + // Count lives after constructor + vector constructor. We only mutate the declared count: the + // preflight must reject before generated Decode attempts a long loop/allocation. + encoded.Buf[8]++ + if err := validateFirstVectorCount(&encoded, maxServiceMessageIDs); err == nil { + t.Fatal("service vector above cap unexpectedly accepted") + } +} + +func TestContainerDecodeUsesBudgetedZeroCopyBodies(t *testing.T) { + encoded := bin.Buffer{} + wantBody := []byte{0x11, 0x22, 0x33, 0x44} + message := proto.Message{ID: 1, SeqNo: 1, Bytes: len(wantBody), Body: wantBody} + if err := (&proto.MessageContainer{Messages: []proto.Message{message}}).Encode(&encoded); err != nil { + t.Fatalf("encode container: %v", err) + } + + s := New(Options{Logger: zaptest.NewLogger(t)}) + s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes - 1) + if _, release, err := s.decodeMessageContainerViews(&encoded, 1); err == nil { + release() + t.Fatal("descriptor allocation unexpectedly bypassed process budget") + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("failed descriptor reservation leaked %d bytes", got) + } + + s.frameBudget = newInboundFrameBudget(2 * containerDescriptorBudgetBytes) + container, release, err := s.decodeMessageContainerViews(&encoded, 1) + if err != nil { + t.Fatalf("decode budgeted container: %v", err) + } + if got := s.frameBudget.usedBytes(); got != containerDescriptorBudgetBytes { + t.Fatalf("descriptor budget = %d, want %d", got, containerDescriptorBudgetBytes) + } + container.Messages[0].Body[0] = 0x99 + if encoded.Buf[8+16] != 0x99 { + t.Fatal("container body was copied instead of viewing the charged input frame") + } + release() + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("released descriptor budget = %d, want zero", got) + } + + truncated := bin.Buffer{Buf: encoded.Buf[:len(encoded.Buf)-1]} + if _, _, err := s.decodeMessageContainerViews(&truncated, 1); err == nil { + t.Fatal("truncated container unexpectedly decoded") + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("failed container decode leaked %d bytes", got) + } +} + +func TestServiceInfoViewsRejectOversizedBytesWithoutDecodeCopy(t *testing.T) { + state := mt.MsgsStateInfo{ReqMsgID: 7, Info: make([]byte, maxServiceMessageIDs)} + var encodedState bin.Buffer + if err := state.Encode(&encodedState); err != nil { + t.Fatalf("encode msgs_state_info: %v", err) + } + reqMsgID, info, err := msgsStateInfoView(&encodedState) + if err != nil || reqMsgID != state.ReqMsgID || len(info) != maxServiceMessageIDs { + t.Fatalf("state info view = id %d len %d err %v", reqMsgID, len(info), err) + } + info[0] = 0x7f + if encodedState.Buf[16] != 0x7f { + t.Fatal("msgs_state_info view unexpectedly copied info") + } + + state.Info = make([]byte, maxServiceMessageIDs+1) + encodedState.Reset() + if err := state.Encode(&encodedState); err != nil { + t.Fatalf("encode oversized msgs_state_info: %v", err) + } + if _, _, err := msgsStateInfoView(&encodedState); err == nil || !strings.Contains(err.Error(), "exceeds") { + t.Fatalf("oversized msgs_state_info err = %v, want capped rejection", err) + } + + all := mt.MsgsAllInfo{MsgIDs: []int64{1, 2}, Info: []byte{4, 4}} + var encodedAll bin.Buffer + if err := all.Encode(&encodedAll); err != nil { + t.Fatalf("encode msgs_all_info: %v", err) + } + count, allInfo, err := msgsAllInfoView(&encodedAll) + if err != nil || count != 2 || len(allInfo) != 2 { + t.Fatalf("all info view = count %d len %d err %v", count, len(allInfo), err) + } + all.Info = make([]byte, maxServiceMessageIDs+1) + encodedAll.Reset() + if err := all.Encode(&encodedAll); err != nil { + t.Fatalf("encode oversized msgs_all_info: %v", err) + } + if _, _, err := msgsAllInfoView(&encodedAll); err == nil || !strings.Contains(err.Error(), "exceeds") { + t.Fatalf("oversized msgs_all_info err = %v, want capped rejection", err) + } +} + +func TestDispatchRejectsExcessiveWrapperDepthBeforeRPC(t *testing.T) { + var body bin.Buffer + if err := (&mt.MsgsStateInfo{ReqMsgID: 1, Info: []byte{4}}).Encode(&body); err != nil { + t.Fatalf("encode leaf: %v", err) + } + encoded := body.Copy() + for i := 0; i < maxDispatchDepth+1; i++ { + var wrapped bin.Buffer + if err := (proto.GZIP{Data: encoded}).Encode(&wrapped); err != nil { + t.Fatalf("encode gzip depth %d: %v", i+1, err) + } + encoded = wrapped.Copy() + } + + s := New(Options{Logger: zaptest.NewLogger(t)}) + var acks []int64 + err := s.dispatch(context.Background(), newConnState(), nil, 4, 0, &bin.Buffer{Buf: encoded}, &acks) + if err == nil || !strings.Contains(err.Error(), "wrapper depth") { + t.Fatalf("deep wrapper err = %v, want wrapper depth rejection", err) + } +} + +func TestOversizedConnectionBuffersAreReleasedAfterFrame(t *testing.T) { + inbound := &bin.Buffer{Buf: make([]byte, 1, maxRetainedConnBuffer+1)} + trimOversizedInboundBuffer(inbound) + if inbound.Buf != nil { + t.Fatalf("oversized inbound buffer cap=%d, want released", cap(inbound.Buf)) + } + regular := &bin.Buffer{Buf: make([]byte, 1, maxRetainedConnBuffer)} + trimOversizedInboundBuffer(regular) + if cap(regular.Buf) != maxRetainedConnBuffer { + t.Fatalf("regular inbound buffer cap=%d, want retained", cap(regular.Buf)) + } + + pool := newOutboundScratchPool(16 << 20) + scratch, err := pool.acquire(context.Background(), nil, maxRetainedConnBuffer+1) + if err != nil { + t.Fatalf("acquire oversized outbound scratch: %v", err) + } + pool.release(scratch) + if got := pool.snapshot(); got != 0 { + t.Fatalf("oversized outbound scratch retained %d bytes, want 0", got) + } +} + +func TestGZIPExpansionUsesProcessBudgetBeforeDecode(t *testing.T) { + payload := make([]byte, 1<<20) + var wrapped bin.Buffer + if err := (proto.GZIP{Data: payload}).Encode(&wrapped); err != nil { + t.Fatalf("encode gzip: %v", err) + } + + s := New(Options{Logger: zaptest.NewLogger(t)}) + s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes - 1) + if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil { + release() + t.Fatal("gzip decode unexpectedly bypassed saturated process budget") + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("failed gzip reservation leaked %d bytes", got) + } + + s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes) + decoded, release, err := s.decodeGZIPWithGlobalBudget(&wrapped) + if err != nil { + t.Fatalf("budgeted gzip decode: %v", err) + } + if len(decoded) != len(payload) { + t.Fatalf("decoded bytes = %d, want %d", len(decoded), len(payload)) + } + if got := s.frameBudget.usedBytes(); got != int64(len(payload)) { + t.Fatalf("held expansion budget = %d, want %d", got, len(payload)) + } + release() + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("released expansion budget = %d, want zero", got) + } +} diff --git a/internal/mtprotoedge/transport_compat.go b/internal/mtprotoedge/transport_compat.go index 7353b78e..5cfa87e0 100644 --- a/internal/mtprotoedge/transport_compat.go +++ b/internal/mtprotoedge/transport_compat.go @@ -34,16 +34,21 @@ type quickAckTransport interface { SendQuickAck(ctx context.Context, token uint32) error } +type deadlineQuickAckTransport interface { + SendQuickAckDeadline(deadline time.Time, token uint32) error +} + type compatTransportListener struct { codec func() transport.Codec listener net.Listener + budget *inboundFrameBudget } -func newCompatTransportListener(codec func() transport.Codec, listener net.Listener) transportListener { - if codec != nil { - return transport.ListenCodec(codec, listener) +func newCompatTransportListener(codec func() transport.Codec, listener net.Listener, budget *inboundFrameBudget) transportListener { + if budget == nil { + panic("mtprotoedge: nil inbound frame budget") } - return &compatTransportListener{listener: listener} + return &compatTransportListener{codec: codec, listener: listener, budget: budget} } // singleConnListener 是一个只产出一条「已接受」连接、随后阻塞到关闭的 net.Listener。 @@ -89,9 +94,27 @@ func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) { } }() - connCodec, reader, err := detectCompatCodec(conn) - if err != nil { - return nil, errors.Wrap(err, "detect codec") + var ( + connCodec transport.Codec + reader io.Reader = conn + ) + if l.codec != nil { + connCodec = l.codec() + if classifyInboundFrameCodec(connCodec) == inboundFrameCodecUnknown { + // Unknown codecs are rejected before their header or first frame is read. Without an + // explicit preflight contract, calling Codec.Read could allocate from an attacker- + // controlled length before the process-wide budget can be reserved. + return nil, errInboundFrameCodecUnsupported + } + if err := connCodec.ReadHeader(conn); err != nil { + return nil, errors.Wrap(err, "read codec header") + } + } else { + var err error + connCodec, reader, err = detectCompatCodec(conn) + if err != nil { + return nil, errors.Wrap(err, "detect codec") + } } return &compatTransportConn{ @@ -99,7 +122,8 @@ func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) { reader: reader, Conn: conn, }, - codec: connCodec, + codec: connCodec, + budget: l.budget, }, nil } @@ -121,11 +145,17 @@ func (w wrappedCompatConn) Read(p []byte) (int, error) { } type compatTransportConn struct { - conn net.Conn - codec transport.Codec + conn net.Conn + codec transport.Codec + budget *inboundFrameBudget readMux sync.Mutex writeMux sync.Mutex + + frameMu sync.Mutex + heldFrameBytes int64 + frameDelivered bool + closed bool } func (c *compatTransportConn) Send(ctx context.Context, b *bin.Buffer) error { @@ -157,6 +187,11 @@ func (c *compatTransportConn) ConsumeQuickAckRequested() bool { } func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) error { + deadline, _ := ctx.Deadline() + return c.SendQuickAckDeadline(deadline, token) +} + +func (c *compatTransportConn) SendQuickAckDeadline(deadline time.Time, token uint32) error { q, ok := c.codec.(quickAckCodec) if !ok { return nil @@ -165,7 +200,6 @@ func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) er c.writeMux.Lock() defer c.writeMux.Unlock() - deadline, _ := ctx.Deadline() if err := c.conn.SetWriteDeadline(deadline); err != nil { return errors.Wrap(err, "set write deadline") } @@ -189,19 +223,185 @@ func (c *compatTransportConn) RecvDeadline(deadline time.Time, b *bin.Buffer) er c.readMux.Lock() defer c.readMux.Unlock() + // Starting the next Recv proves the previous frame slices are no longer consumed, but its + // reusable backing remains live. Keep the high-water reservation until the new length prefix + // atomically grows/reuses it; serveConn later shrinks it to the actually retained capacities. + c.beginInboundFrameRead() if err := c.conn.SetReadDeadline(deadline); err != nil { + c.releaseInboundFrame() return errors.Wrap(err, "set read deadline") } - if err := c.codec.Read(c.conn, b); err != nil { + if err := c.readInboundFrame(b); err != nil { + // A short payload or protocol error cannot escape while retaining a reservation. + c.releaseInboundFrame() return errors.Wrap(err, "read") } return nil } func (c *compatTransportConn) Close() error { + c.frameMu.Lock() + c.closed = true + // Never release here: Close can race both a delivered frame owned by serveConn and a codec read + // still writing into b. Recv's error path or serveConn's deferred ownership release is the + // unique point where those backings become dead. + c.frameMu.Unlock() return c.conn.Close() } +func (c *compatTransportConn) readInboundFrame(b *bin.Buffer) error { + kind := classifyInboundFrameCodec(c.codec) + if kind == inboundFrameCodecUnknown { + return errInboundFrameCodecUnsupported + } + reserveCalls := 0 + reserved := false + var reserveErr error + reserve := func(wireBytes, plaintextBytes int64) error { + reserveCalls++ + if reserveCalls != 1 { + reserveErr = errors.New("inbound frame codec reserved more than once") + return reserveErr + } + reserveErr = c.reserveInboundFrame(wireBytes, plaintextBytes) + reserved = reserveErr == nil + return reserveErr + } + + var err error + if kind == inboundFrameCodecCustom { + custom := unwrapInboundFrameBudgetedCodec(c.codec) + if custom == nil { + return errInboundFrameCodecUnsupported + } + err = custom.ReadWithInboundFrameBudget(c.conn, b, reserve) + } else { + preflight := &inboundFramePreflightReader{r: c.conn, kind: kind, reserve: reserve} + err = c.codec.Read(preflight, b) + } + if err != nil { + return err + } + if reserveErr != nil { + return reserveErr + } + if reserveCalls != 1 || !reserved { + return errInboundFrameNotReserved + } + return c.markInboundFrameDelivered() +} + +func (c *compatTransportConn) reserveInboundFrame(wireBytes, plaintextBytes int64) error { + c.frameMu.Lock() + defer c.frameMu.Unlock() + if c.closed { + return net.ErrClosed + } + n, err := c.budget.growReservation(c.heldFrameBytes, wireBytes, plaintextBytes) + if err != nil { + return err + } + c.heldFrameBytes = n + return nil +} + +func (c *compatTransportConn) beginInboundFrameRead() { + c.frameMu.Lock() + c.frameDelivered = false + c.frameMu.Unlock() +} + +// retainInboundFrameBytes shrinks the high-water frame charge to the capacities that serveConn +// intentionally keeps for reuse after dispatch. It may grow only to account allocator rounding; +// callers drop both buffers and retry with zero when that extra admission is unavailable. +func (c *compatTransportConn) retainInboundFrameBytes(n int64) bool { + if n < 0 { + return false + } + c.frameMu.Lock() + old := c.heldFrameBytes + if n > old { + grown, err := c.budget.growReservation(old, n, 0) + if err != nil { + c.frameMu.Unlock() + return false + } + c.heldFrameBytes = grown + c.frameMu.Unlock() + return true + } + c.heldFrameBytes = n + c.frameMu.Unlock() + c.budget.release(old - n) + return true +} + +func (c *compatTransportConn) releaseInboundFrame() { + c.frameMu.Lock() + n := c.heldFrameBytes + c.heldFrameBytes = 0 + c.frameDelivered = false + c.frameMu.Unlock() + c.budget.release(n) +} + +func (c *compatTransportConn) markInboundFrameDelivered() error { + c.frameMu.Lock() + defer c.frameMu.Unlock() + if c.closed { + // Do not hand a frame to the consumer after Close. The read error path keeps ownership + // accounting until the codec has stopped touching its backing, then releases it. + return net.ErrClosed + } + if c.heldFrameBytes == 0 { + return errInboundFrameNotReserved + } + c.frameDelivered = true + return nil +} + +type inboundFrameOwnershipReleaser interface { + releaseInboundFrame() +} + +type inboundFrameBackingRetainer interface { + retainInboundFrameBytes(int64) bool +} + +func releaseInboundFrameOwnership(conn transport.Conn) { + if releaser, ok := conn.(inboundFrameOwnershipReleaser); ok { + releaser.releaseInboundFrame() + } +} + +// retainInboundFrameBackings transfers the current-frame reservation into a persistent charge +// for reusable buffer capacities. If allocator rounding would exceed the available budget, drop +// both backings and release the reservation rather than retaining unaccounted memory. +func retainInboundFrameBackings(conn transport.Conn, buffers ...*bin.Buffer) { + retainer, ok := conn.(inboundFrameBackingRetainer) + if !ok { + return + } + var retained int64 + for _, b := range buffers { + if b == nil { + continue + } + retained += int64(cap(b.Buf)) + } + if retainer.retainInboundFrameBytes(retained) { + return + } + for _, b := range buffers { + if b != nil { + b.Buf = nil + } + } + if !retainer.retainInboundFrameBytes(0) { + panic("mtprotoedge: failed to release inbound frame backing reservation") + } +} + func detectCompatCodec(c io.Reader) (transport.Codec, io.Reader, error) { var buf [4]byte if _, err := io.ReadFull(c, buf[:1]); err != nil { @@ -233,7 +433,6 @@ type quickAckCodec interface { type quickAckAbridgedCodec struct { quickAckRequested bool - wbuf []byte } func (*quickAckAbridgedCodec) WriteHeader(w io.Writer) error { @@ -261,7 +460,7 @@ func (q *quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error { header[3] = byte(words >> 16) headerLen = 4 } - return writeCompatPacket(w, &q.wbuf, header[:headerLen], b.Raw()) + return writeCompatPacket(w, header[:headerLen], b.Raw()) } func (q *quickAckAbridgedCodec) Read(r io.Reader, b *bin.Buffer) error { @@ -287,7 +486,6 @@ func (*quickAckAbridgedCodec) quickAckResponse(token uint32) [4]byte { type quickAckIntermediateCodec struct { quickAckRequested bool - wbuf []byte } func (*quickAckIntermediateCodec) WriteHeader(w io.Writer) error { @@ -304,7 +502,7 @@ func (q *quickAckIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error { } var header [4]byte binary.LittleEndian.PutUint32(header[:], uint32(b.Len())) - return writeCompatPacket(w, &q.wbuf, header[:], b.Raw()) + return writeCompatPacket(w, header[:], b.Raw()) } func (q *quickAckIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error { @@ -330,7 +528,6 @@ func (*quickAckIntermediateCodec) quickAckResponse(token uint32) [4]byte { type quickAckPaddedIntermediateCodec struct { quickAckRequested bool - wbuf []byte rand *bufio.Reader } @@ -356,13 +553,11 @@ func (q *quickAckPaddedIntermediateCodec) Write(w io.Writer, b *bin.Buffer) erro return err } n := int(padding[0] % 4) - // header(4B) + payload + padding 一次拼进复用缓冲，单次 Write 出站。 - buf := append(q.wbuf[:0], 0, 0, 0, 0) - binary.LittleEndian.PutUint32(buf[:4], uint32(b.Len()+n)) - buf = append(buf, b.Raw()...) - buf = append(buf, padding[:n]...) - q.wbuf = buf - return writeAll(w, buf) + var header [4]byte + binary.LittleEndian.PutUint32(header[:], uint32(b.Len()+n)) + buffers := net.Buffers{header[:], b.Raw(), padding[:n]} + _, err := buffers.WriteTo(w) + return err } func (q *quickAckPaddedIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error { @@ -449,13 +644,13 @@ func validateOutgoingCompatMessage(b *bin.Buffer) error { return nil } -// writeCompatPacket 把 header+payload 拼进调用方持有的复用缓冲后单次写出： -// 保持 MTProto 帧单包出站（quick ack 尾延迟契约），同时避免每帧分配拼包缓冲。 -func writeCompatPacket(w io.Writer, scratch *[]byte, header, payload []byte) error { - buf := append((*scratch)[:0], header...) - buf = append(buf, payload...) - *scratch = buf - return writeAll(w, buf) +// writeCompatPacket avoids a full-frame codec copy. net.Buffers uses vectored I/O for raw TCP +// (one syscall); wrapped writers may receive ordered writes, still serialized by writeMux. The +// outbound scratch lease keeps the encrypted payload alive until all segments finish. +func writeCompatPacket(w io.Writer, header, payload []byte) error { + buffers := net.Buffers{header, payload} + _, err := buffers.WriteTo(w) + return err } func writeAll(w io.Writer, p []byte) error { diff --git a/internal/mtprotoedge/transport_compat_test.go b/internal/mtprotoedge/transport_compat_test.go index c863021c..e8ca2cf8 100644 --- a/internal/mtprotoedge/transport_compat_test.go +++ b/internal/mtprotoedge/transport_compat_test.go @@ -75,8 +75,8 @@ func TestCompatPaddedIntermediateWriteRoundTrip(t *testing.T) { if err := codec.Write(&out, &payload); err != nil { t.Fatalf("write %d: %v", i, err) } - if out.writes != 1 { - t.Fatalf("write %d: writes = %d, want 1", i, out.writes) + if out.writes < 2 || out.writes > 3 { + t.Fatalf("write %d: writes = %d, want header/payload[/padding] segments", i, out.writes) } total := binary.LittleEndian.Uint32(out.Bytes()[:4]) if int(total) != len(out.Bytes())-4 { @@ -97,7 +97,7 @@ func TestCompatPaddedIntermediateWriteRoundTrip(t *testing.T) { } } -func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) { +func TestCompatTransportCodecsWriteSegmentedPacketWithoutFullCopy(t *testing.T) { var payload bin.Buffer payload.PutInt32(0x01020304) payload.PutInt32(0x05060708) @@ -107,8 +107,8 @@ func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) { if err := (&quickAckAbridgedCodec{}).Write(&out, &payload); err != nil { t.Fatalf("write: %v", err) } - if out.writes != 1 { - t.Fatalf("writes = %d, want 1", out.writes) + if out.writes != 2 { + t.Fatalf("generic writer calls = %d, want header+payload; raw TCP uses vectored I/O", out.writes) } if got, want := out.Bytes()[0], byte(payload.Len()/bin.Word); got != want { t.Fatalf("abridged header = %#x, want %#x", got, want) @@ -120,8 +120,8 @@ func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) { if err := (&quickAckIntermediateCodec{}).Write(&out, &payload); err != nil { t.Fatalf("write: %v", err) } - if out.writes != 1 { - t.Fatalf("writes = %d, want 1", out.writes) + if out.writes != 2 { + t.Fatalf("generic writer calls = %d, want header+payload; raw TCP uses vectored I/O", out.writes) } if got, want := binary.LittleEndian.Uint32(out.Bytes()[:4]), uint32(payload.Len()); got != want { t.Fatalf("intermediate length = %d, want %d", got, want) diff --git a/internal/rpc/account.go b/internal/rpc/account.go index 9c218e6b..315ab21d 100644 --- a/internal/rpc/account.go +++ b/internal/rpc/account.go @@ -490,8 +490,8 @@ func (r *Router) onAccountVerifyEmail(ctx context.Context, req *tg.AccountVerify SentCode: tgEmailSentCode(p.PhoneCodeHash, domain.MaskEmail(email), len(strings.TrimSpace(code))), }, nil } - u, loginMessage, needSignUp, signInErr := r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), p.PhoneNumber, p.PhoneCodeHash, code) - authorization, err := r.finishAuthSignIn(ctx, u, loginMessage, needSignUp, signInErr) + u, _, needSignUp, signInErr := r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), p.PhoneNumber, p.PhoneCodeHash, code) + authorization, err := r.finishAuthSignIn(ctx, u, needSignUp, signInErr) if err != nil { return nil, err } diff --git a/internal/rpc/account_business.go b/internal/rpc/account_business.go index 6c3250da..87466de1 100644 --- a/internal/rpc/account_business.go +++ b/internal/rpc/account_business.go @@ -490,7 +490,7 @@ func (r *Router) recordConnectedBusinessPeerSettings(ctx context.Context, userID } authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, settings, sessionID) + event, _, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, settings, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return err } diff --git a/internal/rpc/account_phone.go b/internal/rpc/account_phone.go index 03e71113..c1d6b296 100644 --- a/internal/rpc/account_phone.go +++ b/internal/rpc/account_phone.go @@ -51,10 +51,12 @@ func (r *Router) onAccountChangePhone(ctx context.Context, req *tg.AccountChange return nil, authKeyUnregisteredErr() } sessionID, _ := SessionIDFrom(ctx) + originRawAuthKeyID := rawAuthKeyIDForOrigin(ctx) result, err := r.deps.Account.ChangePhone( ctx, userID, authKeyID, + originRawAuthKeyID, sessionID, req.PhoneNumber, req.PhoneCodeHash, diff --git a/internal/rpc/account_themes_compat_test.go b/internal/rpc/account_themes_compat_test.go index 058541cd..10dd319a 100644 --- a/internal/rpc/account_themes_compat_test.go +++ b/internal/rpc/account_themes_compat_test.go @@ -6,31 +6,39 @@ import ( "github.com/gotd/td/bin" "github.com/gotd/td/tg" + "github.com/gotd/td/tgerr" + "go.uber.org/zap" + "go.uber.org/zap/zaptest/observer" "telesrv/internal/domain" ) -// TestLegacyThemeWireDecode 验证按 DrKLO 12.8.1 的 theme 构造器(比 gotd schema 新) -// 手写解码后能正确复用现有 handler。直接构造 DrKLO 的 wire 字节喂给 fallback compat 层。 -func TestLegacyThemeWireDecode(t *testing.T) { +// TestLegacyThemeWireDispatch 验证 DrKLO theme 构造器经过 Router.Dispatch 的完整链路： +// layerwire 结构预检 -> gotd dispatcher fallback -> compat 解码 -> 现有 handler。 +// 不能只直调 tryLegacyThemeRPC，否则会掩盖预检早于 fallback 的路由回归。 +func TestLegacyThemeWireDispatch(t *testing.T) { const userID = 1000010 ctx := WithUserID(context.Background(), userID) + var authKeyID [8]byte + authKeyID[0] = 1 + const sessionID = 99 files := &fakeFiles{docs: map[int64]domain.Document{ 777: {ID: 777, AccessHash: 7, DCID: 2, MimeType: "application/x-tgtheme-android", Size: 4096}, }} r := newThemeRouter(t, files) - // createTheme 0x8432c21f:flags(=4,document) + slug + title + InputDocument。 + // createTheme 0x8432c21f:flags(document+single settings) + slug + title。 var cb bin.Buffer cb.PutID(legacyCreateThemeID) - cb.PutInt32(1 << 2) // document present - cb.PutString("") // empty slug → auto + cb.PutInt32((1 << 2) | (1 << 3)) + cb.PutString("") // empty slug → auto cb.PutString("Legacy Theme") (&tg.InputDocument{ID: 777, AccessHash: 7}).Encode(&cb) + (&tg.InputThemeSettings{BaseTheme: &tg.BaseThemeDay{}, AccentColor: 0x3997d3}).Encode(&cb) - enc, handled, err := r.tryLegacyThemeRPC(ctx, &cb) - if !handled || err != nil { - t.Fatalf("createTheme legacy = handled %v err %v", handled, err) + enc, err := r.Dispatch(ctx, authKeyID, sessionID, &cb) + if err != nil { + t.Fatalf("createTheme legacy dispatch: %v", err) } th, ok := enc.(*tg.Theme) if !ok { @@ -44,9 +52,29 @@ func TestLegacyThemeWireDecode(t *testing.T) { } else if d, _ := doc.(*tg.Document); d == nil || d.ID != 777 { t.Fatalf("created theme document = %#v, want id 777", doc) } + if settings, ok := th.GetSettings(); !ok || len(settings) != 1 || settings[0].AccentColor != 0x3997d3 { + t.Fatalf("created theme settings = %#v ok=%v, want one legacy setting", settings, ok) + } mustEncodeTheme(t, th) slug := th.Slug + // updateTheme 0x5cb367d5:flags(=2,title) + format + InputTheme + title。 + var ub bin.Buffer + ub.PutID(legacyUpdateThemeID) + ub.PutInt32(1 << 1) + ub.PutString("android") + (&tg.InputTheme{ID: th.ID, AccessHash: th.AccessHash}).Encode(&ub) + ub.PutString("Legacy Theme Updated") + + enc, err = r.Dispatch(ctx, authKeyID, sessionID, &ub) + if err != nil { + t.Fatalf("updateTheme legacy dispatch: %v", err) + } + updated, ok := enc.(*tg.Theme) + if !ok || updated.Title != "Legacy Theme Updated" { + t.Fatalf("updateTheme legacy result = %#v, want updated title", enc) + } + // getTheme 0x8d9d742b:format + InputThemeSlug + document_id(被忽略)。 var gb bin.Buffer gb.PutID(legacyGetThemeID) @@ -54,9 +82,9 @@ func TestLegacyThemeWireDecode(t *testing.T) { (&tg.InputThemeSlug{Slug: slug}).Encode(&gb) gb.PutLong(12345) // document_id ignored - enc, handled, err = r.tryLegacyThemeRPC(ctx, &gb) - if !handled || err != nil { - t.Fatalf("getTheme legacy = handled %v err %v", handled, err) + enc, err = r.Dispatch(ctx, authKeyID, sessionID, &gb) + if err != nil { + t.Fatalf("getTheme legacy dispatch: %v", err) } got, ok := enc.(*tg.Theme) if !ok || got.Slug != slug { @@ -73,18 +101,44 @@ func TestLegacyThemeWireDecode(t *testing.T) { ib.PutString("android") (&tg.InputThemeSlug{Slug: slug}).Encode(&ib) - enc, handled, err = r.tryLegacyThemeRPC(ctx, &ib) - if !handled || err != nil { - t.Fatalf("installTheme legacy = handled %v err %v", handled, err) + enc, err = r.Dispatch(ctx, authKeyID, sessionID, &ib) + if err != nil { + t.Fatalf("installTheme legacy dispatch: %v", err) } if _, ok := enc.(*tg.BoolTrue); !ok { t.Fatalf("installTheme legacy result = %T, want *tg.BoolTrue", enc) } - // 非 theme 构造器 → 不处理。 - var ob bin.Buffer - ob.PutID(0x12345678) - if _, handled, _ := r.tryLegacyThemeRPC(ctx, &ob); handled { - t.Fatalf("unrelated ctor should not be handled") + // 已声明 legacy 方法仍必须精确消费完整结构，截断字段不能到手写 decoder。 + var malformed bin.Buffer + malformed.PutID(legacyCreateThemeID) + malformed.PutInt32(0) + malformed.PutString("slug") // missing title + if _, err := r.Dispatch(ctx, authKeyID, sessionID, &malformed); !tgerr.Is(err, "INPUT_REQUEST_INVALID") { + t.Fatalf("malformed legacy theme err = %v, want INPUT_REQUEST_INVALID", err) + } +} + +func TestUnknownRPCReachesCompatibilityTraceAfterOpaquePreflight(t *testing.T) { + const unknownID = uint32(0x12345678) + r := newThemeRouter(t, &fakeFiles{}) + r.deps.Auth = &captureAuthService{} + core, logs := observer.New(zap.WarnLevel) + r.log = zap.New(core) + + var b bin.Buffer + b.PutID(unknownID) + // Deliberately resembles a forged vector count. Because the constructor is unknown, the + // body remains opaque and is never decoded or allocated from; total frame/RPC budgets bound it. + b.PutUint32(0xffffffff) + if _, err := r.Dispatch(context.Background(), [8]byte{1}, 101, &b); !tgerr.Is(err, "NOT_IMPLEMENTED") { + t.Fatalf("unknown dispatch err = %v, want NOT_IMPLEMENTED", err) + } + entries := logs.FilterMessage("Unhandled RPC (compatibility trace)").All() + if len(entries) != 1 { + t.Fatalf("compatibility trace entries = %d, want 1", len(entries)) + } + if got, ok := entries[0].ContextMap()["type_id"]; !ok || got != "0x12345678" { + t.Fatalf("trace type_id = %#v, want %#x", got, unknownID) } } diff --git a/internal/rpc/album_group.go b/internal/rpc/album_group.go new file mode 100644 index 00000000..546527cd --- /dev/null +++ b/internal/rpc/album_group.go @@ -0,0 +1,39 @@ +package rpc + +import ( + "context" + "errors" + + "go.uber.org/zap" + + "telesrv/internal/domain" +) + +func (r *Router) reserveAlbumGroup(ctx context.Context, userID int64, peer domain.Peer, items []domain.AlbumGroupReservationItem) (int64, error) { + reservations, ok := r.deps.Messages.(AlbumGroupService) + if !ok { + r.log.Error("messages.sendMultiMedia album reservation capability missing", + append(r.contextLogFields(ctx), zap.Int64("user_id", userID), zap.String("peer_type", string(peer.Type)), zap.Int64("peer_id", peer.ID))...) + return 0, internalErr() + } + groupedID, err := reservations.ReserveAlbumGroup(ctx, userID, domain.AlbumGroupReservationRequest{ + SenderUserID: userID, + Peer: peer, + Items: items, + ProposedGroupedID: randomNonZeroInt64(), + }) + if errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + return 0, randomIDDuplicateErr() + } + if err != nil { + r.log.Error("messages.sendMultiMedia album reservation failed", + append(r.contextLogFields(ctx), zap.Error(err), zap.Int64("user_id", userID), zap.String("peer_type", string(peer.Type)), zap.Int64("peer_id", peer.ID), zap.Int("items", len(items)))...) + return 0, internalErr() + } + if groupedID == 0 { + r.log.Error("messages.sendMultiMedia album reservation returned zero grouped_id", + append(r.contextLogFields(ctx), zap.Int64("user_id", userID), zap.String("peer_type", string(peer.Type)), zap.Int64("peer_id", peer.ID))...) + return 0, internalErr() + } + return groupedID, nil +} diff --git a/internal/rpc/auth.go b/internal/rpc/auth.go index 5e84417e..fbda59ac 100644 --- a/internal/rpc/auth.go +++ b/internal/rpc/auth.go @@ -241,6 +241,9 @@ func (r *Router) pushLoginTokenAccepted(ctx context.Context, target loginTokenTa // 若该手机号账号设置了登录邮箱，验证码改投递到邮箱，返回 sentCodeTypeEmailCode // （客户端据此进入"输入邮箱验证码"界面，随后用 auth.signIn 的 email_verification 完成登录）。 func (r *Router) onAuthSendCode(ctx context.Context, req *tg.AuthSendCodeRequest) (tg.AuthSentCodeClass, error) { + if err := r.checkAuthCodeRateLimit(ctx, req.PhoneNumber); err != nil { + return nil, err + } r.rememberClientAPIID(ctx, req.APIID) hash, err := r.deps.Auth.SendCode(ctx, req.PhoneNumber) if err != nil { @@ -328,25 +331,21 @@ func (r *Router) tgSentCodeForHash(ctx context.Context, hash string) (tg.AuthSen // 带 email_verification 时走登录邮箱路径（验证码来自邮箱而非短信）。 func (r *Router) onAuthSignIn(ctx context.Context, req *tg.AuthSignInRequest) (tg.AuthAuthorizationClass, error) { var ( - u domain.User - loginMessage domain.Message - needSignUp bool - err error + u domain.User + needSignUp bool + err error ) if verification, ok := req.GetEmailVerification(); ok { - u, loginMessage, needSignUp, err = r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, emailVerificationCode(verification)) + u, _, needSignUp, err = r.deps.Auth.SignInWithEmail(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, emailVerificationCode(verification)) } else { - u, loginMessage, needSignUp, err = r.deps.Auth.SignIn(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, req.PhoneCode) + u, _, needSignUp, err = r.deps.Auth.SignIn(ctx, r.authzFromCtx(ctx), req.PhoneNumber, req.PhoneCodeHash, req.PhoneCode) } - return r.finishAuthSignIn(ctx, u, loginMessage, needSignUp, err) + return r.finishAuthSignIn(ctx, u, needSignUp, err) } -func (r *Router) finishAuthSignIn(ctx context.Context, u domain.User, loginMessage domain.Message, needSignUp bool, err error) (tg.AuthAuthorizationClass, error) { +func (r *Router) finishAuthSignIn(ctx context.Context, u domain.User, needSignUp bool, err error) (tg.AuthAuthorizationClass, error) { if err != nil { if errors.Is(err, domain.ErrSessionPasswordNeeded) && u.ID != 0 { - if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil { - return nil, internalErr() - } // 两步验证未完成：绝不能把 auth_key/session 标记为已登录，否则客户端忽略 // SESSION_PASSWORD_NEEDED、直接调用业务 RPC 即可绕过 2FA。失效缓存并把 session // 置为未授权，让后续鉴权重新读到 password_pending 并拒绝；待 checkPassword 通过后再授权。 @@ -360,19 +359,18 @@ func (r *Router) finishAuthSignIn(ctx context.Context, u domain.User, loginMessa if needSignUp { return &tg.AuthAuthorizationSignUpRequired{}, nil } - if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil { - return nil, internalErr() - } if id, ok := AuthKeyIDFrom(ctx); ok { r.setAuthUserCache(id, u.ID, true) } r.bindSessionUser(ctx, u.ID) - r.enqueueLoginMessageBootstrap(ctx, loginMessage) r.pushSignInServiceNotificationToOthers(ctx, u) return &tg.AuthAuthorization{User: r.tgSelfUser(u)}, nil } func (r *Router) onAuthResendCode(ctx context.Context, req *tg.AuthResendCodeRequest) (tg.AuthSentCodeClass, error) { + if err := r.checkAuthCodeRateLimit(ctx, req.PhoneNumber); err != nil { + return nil, err + } var hash string var err error if scoped, ok := r.deps.Auth.(interface { @@ -545,15 +543,34 @@ func (r *Router) completePendingPasswordSignIn(ctx context.Context, authKeyID [8 // onAuthResetLoginEmail 处理 auth.resetLoginEmail：用户登录设备时无法访问登录邮箱时 // 清除登录邮箱，改回手机验证码登录，返回一个新的手机 sentCode 供其继续。 +type loginEmailResetConsumer interface { + ConsumeLoginEmailReset(ctx context.Context, phone, phoneCodeHash string) (userID int64, err error) + SendPhoneCodeAfterLoginEmailReset(ctx context.Context, phone string, expectedUserID int64) (string, error) +} + func (r *Router) onAuthResetLoginEmail(ctx context.Context, req *tg.AuthResetLoginEmailRequest) (tg.AuthSentCodeClass, error) { if r.deps.Account == nil || r.deps.Auth == nil { return nil, internalErr() } - if err := r.deps.Account.ClearLoginEmailByPhone(ctx, req.PhoneNumber); err != nil { + if err := r.checkAuthCodeRateLimit(ctx, req.PhoneNumber); err != nil { + return nil, err + } + resetConsumer, ok := r.deps.Auth.(loginEmailResetConsumer) + if !ok { return nil, internalErr() } - hash, err := r.deps.Auth.SendCode(ctx, req.PhoneNumber) + resetUserID, err := resetConsumer.ConsumeLoginEmailReset(ctx, req.PhoneNumber, req.PhoneCodeHash) if err != nil { + return nil, signInErr(err) + } + if err := r.deps.Account.ClearLoginEmail(ctx, resetUserID); err != nil { + return nil, internalErr() + } + hash, err := resetConsumer.SendPhoneCodeAfterLoginEmailReset(ctx, req.PhoneNumber, resetUserID) + if err != nil { + if errors.Is(err, auth.ErrCodeExpired) || errors.Is(err, auth.ErrCodeInvalid) { + return nil, signInErr(err) + } if errors.Is(err, auth.ErrPhoneNumberInvalid) || errors.Is(err, auth.ErrSystemUserLoginForbidden) { return nil, phoneNumberInvalidErr() @@ -564,7 +581,7 @@ func (r *Router) onAuthResetLoginEmail(ctx context.Context, req *tg.AuthResetLog } // emailVerificationCode 从 emailVerification 取出可校验的字符串（验证码 / Google·Apple -// 令牌）。开发环境一律按"任意非空即通过"处理，故三者等价取值。 +// 令牌）；最终必须由 auth service 对签发记录精确校验。 // onAuthInitPasskeyLogin 处理 auth.initPasskeyLogin：生成一次性断言挑战（discoverable）， // 以 DataJSON（顶层含 publicKey）返回。免授权（登录前）。 func (r *Router) onAuthInitPasskeyLogin(ctx context.Context, req *tg.AuthInitPasskeyLoginRequest) (*tg.AuthPasskeyLoginOptions, error) { @@ -579,7 +596,8 @@ func (r *Router) onAuthInitPasskeyLogin(ctx context.Context, req *tg.AuthInitPas } // onAuthFinishPasskeyLogin 处理 auth.finishPasskeyLogin：验证登录断言并绑定 auth_key。 -// 收尾与 signIn 同构（清水位→授权缓存→session 绑定）；passkey 是强因子，直接完全授权 +// 收尾与 signIn 同构（Bind 原子切换 update baseline → 授权缓存 → session 绑定）； +// passkey 是强因子，直接完全授权 // （不走 SESSION_PASSWORD_NEEDED）。FromDCID/FromAuthKeyID 为多 DC 重路由用，本单 DC 忽略。 func (r *Router) onAuthFinishPasskeyLogin(ctx context.Context, req *tg.AuthFinishPasskeyLoginRequest) (tg.AuthAuthorizationClass, error) { if r.deps.Passkey == nil || r.deps.Auth == nil { @@ -597,9 +615,6 @@ func (r *Router) onAuthFinishPasskeyLogin(ctx context.Context, req *tg.AuthFinis if err != nil { return nil, passkeyErr(err) } - if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil { - return nil, internalErr() - } if id, ok := AuthKeyIDFrom(ctx); ok { r.setAuthUserCache(id, u.ID, true) } @@ -621,7 +636,7 @@ func emailVerificationCode(v tg.EmailVerificationClass) string { // onAuthImportBotAuthorization 处理 auth.importBotAuthorization：bot 程序凭 token // 登录为 bot 账号。api_id/api_hash 与现有 sendCode 行为一致不校验（无 app 注册表）。 -// 收尾与 signIn 同构（清水位→授权缓存→session 绑定），但不写登录消息、不推 +// 收尾与 signIn 同构（Bind 原子切换 update baseline → 授权缓存 → session 绑定），但不写登录消息、不推 // signIn 服务通知——那是手机登录语义。 func (r *Router) onAuthImportBotAuthorization(ctx context.Context, req *tg.AuthImportBotAuthorizationRequest) (tg.AuthAuthorizationClass, error) { if r.deps.Auth == nil { @@ -631,9 +646,6 @@ func (r *Router) onAuthImportBotAuthorization(ctx context.Context, req *tg.AuthI if err != nil { return nil, importBotAuthorizationErr(err) } - if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil { - return nil, internalErr() - } if id, ok := AuthKeyIDFrom(ctx); ok { r.setAuthUserCache(id, u.ID, true) } @@ -647,9 +659,6 @@ func (r *Router) onAuthSignUp(ctx context.Context, req *tg.AuthSignUpRequest) (t if err != nil { return nil, signInErr(err) } - if err := r.clearAuthKeyStateOnUserChange(ctx, u.ID); err != nil { - return nil, internalErr() - } if id, ok := AuthKeyIDFrom(ctx); ok { r.setAuthUserCache(id, u.ID, true) } @@ -689,18 +698,6 @@ func (r *Router) onAuthLogOut(ctx context.Context) (*tg.AuthLoggedOut, error) { return &tg.AuthLoggedOut{}, nil } -func (r *Router) clearAuthKeyStateOnUserChange(ctx context.Context, newUserID int64) error { - oldUserID, ok := UserIDFrom(ctx) - if !ok || oldUserID == 0 || oldUserID == newUserID { - return nil - } - id, ok := AuthKeyIDFrom(ctx) - if !ok { - return nil - } - return r.clearAuthKeyState(ctx, id) -} - func (r *Router) clearAuthKeyState(ctx context.Context, authKeyID [8]byte) error { if r.deps.Updates == nil { return nil @@ -770,8 +767,9 @@ func (r *Router) pushSignInServiceNotificationToOthers(ctx context.Context, u do return } authKeyID, hasAuthKeyID := AuthKeyIDFrom(ctx) + rawAuthKeyID, hasRawAuthKeyID := RawAuthKeyIDFrom(ctx) sessionID, hasSessionID := SessionIDFrom(ctx) - if !hasAuthKeyID || !hasSessionID { + if !hasAuthKeyID || !hasRawAuthKeyID || !hasSessionID { return } notification := r.tgSignInServiceNotification(ctx, u, authKeyID) @@ -779,7 +777,7 @@ func (r *Router) pushSignInServiceNotificationToOthers(ctx context.Context, u do pushCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() if scoped, ok := r.scopedSessions(); ok { - if sent, err := scoped.PushToUserExceptAuthKeySession(pushCtx, u.ID, authKeyID, sessionID, proto.MessageFromServer, notification); err != nil { + if sent, err := scoped.PushToUserExceptAuthKeySession(pushCtx, u.ID, rawAuthKeyID, sessionID, proto.MessageFromServer, notification); err != nil { r.log.Debug("push sign-in service notification", zap.Int64("user_id", u.ID), zap.Int("sent", sent), zap.Error(err)) } return diff --git a/internal/rpc/auth_code_rate_limit_test.go b/internal/rpc/auth_code_rate_limit_test.go new file mode 100644 index 00000000..566f1d87 --- /dev/null +++ b/internal/rpc/auth_code_rate_limit_test.go @@ -0,0 +1,248 @@ +package rpc + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "strings" + "testing" + "time" + + "github.com/gotd/td/clock" + "github.com/gotd/td/tg" + "go.uber.org/zap/zaptest" + + "telesrv/internal/app/auth" + "telesrv/internal/domain" +) + +type authCodeRateTestService struct { + *captureAuthService + sendCalls int + resendCalls int + resetCalls int + resetPhone string + resetHash string + resetUserID int64 + resetErr error +} + +func (s *authCodeRateTestService) SendCode(context.Context, string) (string, error) { + s.sendCalls++ + return "send-hash", nil +} + +func (s *authCodeRateTestService) ResendCode(context.Context, string, string) (string, error) { + s.resendCalls++ + return "resend-hash", nil +} + +func (s *authCodeRateTestService) ConsumeLoginEmailReset(_ context.Context, phone, hash string) (int64, error) { + s.resetCalls++ + s.resetPhone = phone + s.resetHash = hash + return s.resetUserID, s.resetErr +} + +func (s *authCodeRateTestService) SendPhoneCodeAfterLoginEmailReset(_ context.Context, _ string, expectedUserID int64) (string, error) { + s.sendCalls++ + if expectedUserID != s.resetUserID { + return "", auth.ErrCodeInvalid + } + return "send-hash", nil +} + +type authCodeRateTestAccount struct { + AccountService + clearCalls int + clearUserID int64 +} + +func (s *authCodeRateTestAccount) ClearLoginEmail(_ context.Context, userID int64) error { + s.clearCalls++ + s.clearUserID = userID + return nil +} + +func TestAuthSendCodeRateLimitUsesOpaquePhoneAndRawAuthKeyKeys(t *testing.T) { + phone := "+1 (555) 123-4567" + rawAuthKeyID := [8]byte{0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef} + limiter := &captureRateLimiter{} + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}} + r := New(Config{ + AuthCodePhoneRateLimit: 5, + AuthCodeAuthKeyRateLimit: 20, + AuthCodeRateWindow: 10 * time.Minute, + }, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System) + + ctx := WithRawAuthKeyID(context.Background(), rawAuthKeyID) + if _, err := r.onAuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 2040}); err != nil { + t.Fatalf("onAuthSendCode: %v", err) + } + if authService.sendCalls != 1 { + t.Fatalf("SendCode calls = %d, want 1", authService.sendCalls) + } + if len(limiter.calls) != 2 { + t.Fatalf("limiter calls = %d, want phone + raw auth key", len(limiter.calls)) + } + digest := sha256.Sum256([]byte(domain.NormalizePhone(phone))) + wantPhoneKey := authCodePhoneRateLimitKeyPrefix + hex.EncodeToString(digest[:]) + wantAuthKey := authCodeAuthKeyRateLimitKeyPrefix + hex.EncodeToString(rawAuthKeyID[:]) + if got := limiter.calls[0]; got.key != wantAuthKey || got.cost != 1 || got.limit != 20 || got.window != 10*time.Minute { + t.Fatalf("auth-key limiter call = %+v", got) + } + if got := limiter.calls[1]; got.key != wantPhoneKey || got.cost != 1 || got.limit != 5 || got.window != 10*time.Minute { + t.Fatalf("phone limiter call = %+v", got) + } + for _, call := range limiter.calls { + if strings.Contains(call.key, domain.NormalizePhone(phone)) || strings.Contains(call.key, phone) { + t.Fatalf("limiter key leaked phone: %q", call.key) + } + } +} + +func TestAuthSendCodePhoneRateLimitPrecedesBusinessLookupAndWrite(t *testing.T) { + limiter := &captureRateLimiter{block: true, retryAfter: 17} + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}} + r := New(Config{ + AuthCodePhoneRateLimit: 5, + AuthCodeRateWindow: time.Minute, + }, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System) + + _, err := r.onAuthSendCode(WithRawAuthKeyID(context.Background(), [8]byte{1}), &tg.AuthSendCodeRequest{PhoneNumber: "+86 188 0000 0000", APIID: 2040}) + if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") || !strings.Contains(err.Error(), "(17)") { + t.Fatalf("sendCode err = %v, want FLOOD_WAIT 17", err) + } + if authService.sendCalls != 0 { + t.Fatalf("SendCode calls = %d, want 0", authService.sendCalls) + } + if len(authService.authKeyClientInfos) != 0 { + t.Fatalf("blocked sendCode persisted client info: %+v", authService.authKeyClientInfos) + } + if len(limiter.calls) != 1 || !strings.HasPrefix(limiter.calls[0].key, authCodePhoneRateLimitKeyPrefix) { + t.Fatalf("limiter calls = %+v, want only phone dimension", limiter.calls) + } +} + +func TestAuthSendCodeRawAuthKeyBlockDoesNotCreatePhoneDimension(t *testing.T) { + limiter := &captureRateLimiter{block: true, retryAfter: 31} + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}} + r := New(Config{ + AuthCodePhoneRateLimit: 5, + AuthCodeAuthKeyRateLimit: 20, + AuthCodeRateWindow: time.Minute, + }, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System) + + rawAuthKeyID := [8]byte{7, 7, 7, 7, 7, 7, 7, 7} + _, err := r.onAuthSendCode(WithRawAuthKeyID(context.Background(), rawAuthKeyID), &tg.AuthSendCodeRequest{PhoneNumber: "15550000001"}) + if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") { + t.Fatalf("sendCode err = %v, want FLOOD_WAIT", err) + } + wantKey := authCodeAuthKeyRateLimitKeyPrefix + hex.EncodeToString(rawAuthKeyID[:]) + if len(limiter.calls) != 1 || limiter.calls[0].key != wantKey { + t.Fatalf("limiter calls = %+v, want only raw auth-key %q", limiter.calls, wantKey) + } + if authService.sendCalls != 0 { + t.Fatalf("SendCode calls = %d, want 0", authService.sendCalls) + } +} + +func TestAuthSendCodeInvalidPhoneCreatesNoLimiterKey(t *testing.T) { + limiter := &captureRateLimiter{} + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}} + r := New(Config{ + AuthCodePhoneRateLimit: 5, + AuthCodeAuthKeyRateLimit: 20, + AuthCodeRateWindow: time.Minute, + }, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System) + + _, err := r.onAuthSendCode(WithRawAuthKeyID(context.Background(), [8]byte{1}), &tg.AuthSendCodeRequest{PhoneNumber: "not-a-phone"}) + if err == nil || !strings.Contains(err.Error(), "PHONE_NUMBER_INVALID") { + t.Fatalf("sendCode err = %v, want PHONE_NUMBER_INVALID", err) + } + if len(limiter.calls) != 0 || authService.sendCalls != 0 { + t.Fatalf("invalid phone limiter/service calls = %d/%d, want 0/0", len(limiter.calls), authService.sendCalls) + } +} + +func TestAuthResendCodeRawAuthKeyRateLimitPrecedesRotation(t *testing.T) { + limiter := &captureRateLimiter{block: true, retryAfter: 23} + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}} + r := New(Config{ + AuthCodeAuthKeyRateLimit: 20, + AuthCodeRateWindow: 2 * time.Minute, + }, Deps{Auth: authService, Limiter: limiter}, zaptest.NewLogger(t), clock.System) + + rawAuthKeyID := [8]byte{9, 8, 7, 6, 5, 4, 3, 2} + _, err := r.onAuthResendCode(WithRawAuthKeyID(context.Background(), rawAuthKeyID), &tg.AuthResendCodeRequest{ + PhoneNumber: "8618800000000", + PhoneCodeHash: "old-hash", + }) + if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") || !strings.Contains(err.Error(), "(23)") { + t.Fatalf("resendCode err = %v, want FLOOD_WAIT 23", err) + } + if authService.resendCalls != 0 { + t.Fatalf("ResendCode calls = %d, want 0", authService.resendCalls) + } + wantKey := authCodeAuthKeyRateLimitKeyPrefix + hex.EncodeToString(rawAuthKeyID[:]) + if len(limiter.calls) != 1 || limiter.calls[0].key != wantKey { + t.Fatalf("limiter calls = %+v, want %q", limiter.calls, wantKey) + } +} + +func TestAuthResetLoginEmailRateLimitPrecedesEmailClear(t *testing.T) { + limiter := &captureRateLimiter{block: true, retryAfter: 29} + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}} + accountService := &authCodeRateTestAccount{} + r := New(Config{AuthCodePhoneRateLimit: 5, AuthCodeRateWindow: time.Minute}, Deps{ + Auth: authService, Account: accountService, Limiter: limiter, + }, zaptest.NewLogger(t), clock.System) + + _, err := r.onAuthResetLoginEmail(context.Background(), &tg.AuthResetLoginEmailRequest{ + PhoneNumber: "8618800000000", + PhoneCodeHash: "email-hash", + }) + if err == nil || !strings.Contains(err.Error(), "FLOOD_WAIT") || !strings.Contains(err.Error(), "(29)") { + t.Fatalf("resetLoginEmail err = %v, want FLOOD_WAIT 29", err) + } + if accountService.clearCalls != 0 || authService.resetCalls != 0 || authService.sendCalls != 0 { + t.Fatalf("side effects reset=%d clear=%d send=%d, want 0/0/0", authService.resetCalls, accountService.clearCalls, authService.sendCalls) + } +} + +func TestAuthResetLoginEmailConsumesHashBeforeClearAndResend(t *testing.T) { + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}, resetUserID: 4242} + accountService := &authCodeRateTestAccount{} + r := New(Config{}, Deps{Auth: authService, Account: accountService}, zaptest.NewLogger(t), clock.System) + req := &tg.AuthResetLoginEmailRequest{PhoneNumber: "+1 555 123 9999", PhoneCodeHash: "email-login-hash"} + + result, err := r.onAuthResetLoginEmail(context.Background(), req) + if err != nil { + t.Fatalf("onAuthResetLoginEmail: %v", err) + } + if authService.resetCalls != 1 || authService.resetPhone != req.PhoneNumber || authService.resetHash != req.PhoneCodeHash || + accountService.clearCalls != 1 || accountService.clearUserID != authService.resetUserID || authService.sendCalls != 1 { + t.Fatalf("calls reset=%d(%q,%q uid=%d) clear=%d(uid=%d) send=%d", authService.resetCalls, authService.resetPhone, authService.resetHash, authService.resetUserID, accountService.clearCalls, accountService.clearUserID, authService.sendCalls) + } + sent, ok := result.(*tg.AuthSentCode) + if !ok || sent.PhoneCodeHash != "send-hash" { + t.Fatalf("result=%T %+v, want sentCode/send-hash", result, result) + } +} + +func TestAuthResetLoginEmailInvalidHashNeverClearsFactor(t *testing.T) { + authService := &authCodeRateTestService{captureAuthService: &captureAuthService{}, resetErr: auth.ErrCodeExpired} + accountService := &authCodeRateTestAccount{} + r := New(Config{}, Deps{Auth: authService, Account: accountService}, zaptest.NewLogger(t), clock.System) + + _, err := r.onAuthResetLoginEmail(context.Background(), &tg.AuthResetLoginEmailRequest{ + PhoneNumber: "15551239999", + PhoneCodeHash: "expired-email-hash", + }) + if err == nil || !strings.Contains(err.Error(), "PHONE_CODE_EXPIRED") { + t.Fatalf("onAuthResetLoginEmail err=%v, want PHONE_CODE_EXPIRED", err) + } + if authService.resetCalls != 1 || accountService.clearCalls != 0 || authService.sendCalls != 0 { + t.Fatalf("calls reset=%d clear=%d send=%d, want 1/0/0", authService.resetCalls, accountService.clearCalls, authService.sendCalls) + } +} diff --git a/internal/rpc/auth_qr_login_test.go b/internal/rpc/auth_qr_login_test.go index eba2df6d..d20ae378 100644 --- a/internal/rpc/auth_qr_login_test.go +++ b/internal/rpc/auth_qr_login_test.go @@ -107,15 +107,16 @@ func TestAuthLoginTokenAcceptedByAndroidBindsTargetSession(t *testing.T) { if snap.sessionID != targetSession || snap.userID != scannerUserID || !snap.userResolved { t.Fatalf("target session snapshot = %+v, want session/user/resolved %d/%d/true", snap, targetSession, scannerUserID) } - if snap.messageType != proto.MessageFromServer { - t.Fatalf("push message type = %v, want MessageFromServer", snap.messageType) - } if !sessions.immediatePushSeen() { t.Fatal("login token update was not pushed through the immediate pre-auth path") } - short, ok := snap.message.(*tg.UpdateShort) + immediateType, immediateMessage := sessions.immediatePushSnapshot() + if immediateType != proto.MessageFromServer { + t.Fatalf("immediate push message type = %v, want MessageFromServer", immediateType) + } + short, ok := immediateMessage.(*tg.UpdateShort) if !ok { - t.Fatalf("push message = %T, want *tg.UpdateShort", snap.message) + t.Fatalf("immediate push message = %T, want *tg.UpdateShort", immediateMessage) } if _, ok := short.Update.(*tg.UpdateLoginToken); !ok { t.Fatalf("pushed update = %T, want *tg.UpdateLoginToken", short.Update) diff --git a/internal/rpc/bots_longtail.go b/internal/rpc/bots_longtail.go index 998862c5..d7b5a22b 100644 --- a/internal/rpc/bots_longtail.go +++ b/internal/rpc/bots_longtail.go @@ -151,7 +151,6 @@ func (r *Router) sendBotAllowedServiceMessageWith(ctx context.Context, userID, b if err != nil { return domain.SendPrivateTextResult{}, err } - authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) return r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{ SenderUserID: userID, @@ -165,7 +164,7 @@ func (r *Router) sendBotAllowedServiceMessageWith(ctx context.Context, userID, b }, }, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) diff --git a/internal/rpc/channel_fanout_dispatcher.go b/internal/rpc/channel_fanout_dispatcher.go index 29ad32fc..7b4e37d8 100644 --- a/internal/rpc/channel_fanout_dispatcher.go +++ b/internal/rpc/channel_fanout_dispatcher.go @@ -2,8 +2,10 @@ package rpc import ( "context" + "fmt" "sync" "sync/atomic" + "time" "github.com/gotd/td/tg" "go.uber.org/zap" @@ -23,12 +25,28 @@ import ( // DrKLO Android ~1.5s 乱序窗口与 TDesktop PtsWaiter 的连续性期望（设计 §10.2）。 // - 单实例 + 无 durable 重投队列 + 同 channel 串行 → 无乱序、无自重复，故 v1 不需要 // per-session at-most-once 双水位（那是 Phase 3 跨实例的事，设计 §9/§10.1）。 -// - 有界队列满时丢弃当前 job 并告警：被丢 recipient 会在该 channel 下一条成功投递的 -// pts 跳变时经 getChannelDifference 收敛（设计约束 B）。 +// - 有界队列满时不静默丢弃恢复触发：真实 payload job 降级为按 channel 合并、只保留 +// 最高 pts 的 UpdateChannelTooLong nudge。每个 shard 独立公平 drain；即使该频道随后 +// 静默，也不依赖“下一条消息”才能触发 getChannelDifference（设计约束 B）。 const ( defaultChannelFanoutShards = 64 - defaultChannelFanoutBuffer = 2048 + // The old 64x2048 channel buffers eagerly retained up to 131k closures (each may capture a + // message batch and ~2k recipients). Keep one small FIFO per ordering shard and enforce a + // process-wide retained-byte budget below. + defaultChannelFanoutBuffer = 64 + defaultChannelFanoutMaxQueuedJobs = 4096 + defaultChannelFanoutMaxQueuedBytes = 256 << 20 + defaultChannelFanoutOverflowPerShard = 256 + defaultChannelNudgeWorkers = 8 + defaultChannelNudgeQueue = 4096 + defaultChannelFanoutRecoverySweepPage = 256 + channelFanoutMinRetainedBytes int64 = 64 << 10 + channelFanoutNudgeRetryMin = time.Millisecond + channelFanoutNudgeRetryMax = 50 * time.Millisecond + channelFanoutRecoveryRetryMin = 10 * time.Millisecond + channelFanoutRecoveryRetryMax = time.Second + defaultChannelFanoutNudgeDeadline = 5 * time.Second ) // channelFanoutBuilder 按 viewer 构建该 viewer 视角的 channel updates。与同步 @@ -37,7 +55,7 @@ const ( type channelFanoutBuilder func(ctx context.Context, viewerUserID int64) *tg.Updates // channelFanoutJob 是一条频道 payload fan-out 任务。Pts 仅用于日志/折叠语义；真值仍是 -// channel_update_events，worker 只做在线投递。originAuthKeyID 是业务视角 auth key +// channel_update_events，worker 只做在线投递。originAuthKeyID 是物理 raw auth key // （与 SessionManager.shouldExcludeSession 的比较侧一致），用于显式排除发起设备——异步 // 执行时请求 ctx 已失效，不能再靠 ctx 派生排除。 type channelFanoutJob struct { @@ -50,6 +68,270 @@ type channelFanoutJob struct { originSessionID int64 prefetch channelFanoutPrefetch build channelFanoutBuilder + // retainedBytes is a conservative reservation for the request-derived closure, result + // snapshots and explicit recipient slice. It is charged before the job enters any queue. + retainedBytes int64 + // queueSeq 是 dispatcher shard 内部的 FIFO 序号。仅成功进入正常 payload queue 的 + // job 占用序号；overflow watermark 记录入队失败时已经接受的最大序号，等这些更早 + // payload 处理完后才发 nudge，避免 nudge 越过其之前的正常 FIFO payload。 + queueSeq uint64 +} + +// channelFanoutOverflow 是 queue full 时的 nudge-only 恢复水位。同一 channel 只保留 +// 最大 pts；barrier 是该次 overflow 之前已经进入正常 FIFO 的最后一个 shard 序号。 +type channelFanoutOverflow struct { + pts int + barrier uint64 +} + +// channelFanoutShard 把正常 payload FIFO 与 overflow nudge mailbox 放在同一个 worker +// 下。overflowOrder 每个 channel 最多出现一次；热点 channel 只更新 map 水位，不会占满 +// order，从而不能把其它 channel 的唯一恢复 nudge 永久饿死。 +type channelFanoutShard struct { + jobs chan channelFanoutJob + overflowWake chan struct{} + // overflowSpace is a generation channel, not a one-token notification. A slot + // release closes the current generation and installs a fresh channel while mu is + // held, waking every waiter that observed the old full mailbox. Each waiter then + // competes under mu for the actually available slots; losers observe the new + // generation and sleep again. This avoids losing N-1 wakeups when several slots + // are released before any of N waiters gets scheduled. + overflowSpace chan struct{} + + mu sync.Mutex + nextSeq uint64 + processedSeq uint64 + overflow map[int64]channelFanoutOverflow + overflowOrder []int64 + overflowLimit int + // overflowWaiters is guarded by mu and only covers the distinct-channel + // saturation slow path. Besides making the wait lifecycle explicit, it avoids + // allocating a fresh generation channel when no goroutine is subscribed. + overflowWaiters int +} + +func newChannelFanoutShard(buffer int) *channelFanoutShard { + return &channelFanoutShard{ + jobs: make(chan channelFanoutJob, buffer), + overflowWake: make(chan struct{}, 1), + overflowSpace: make(chan struct{}), + overflow: make(map[int64]channelFanoutOverflow), + overflowLimit: defaultChannelFanoutOverflowPerShard, + } +} + +// enqueue 尝试把真实 payload 放入正常 FIFO；满时按 channel 合并最高 pts 的 nudge-only +// watermark。返回 true 表示正常入队，false 表示已经安全降级为 overflow watermark。 +func (s *channelFanoutShard) enqueue(job channelFanoutJob) bool { + s.mu.Lock() + job.queueSeq = s.nextSeq + 1 + select { + case s.jobs <- job: + s.nextSeq = job.queueSeq + s.mu.Unlock() + return true + default: + s.mu.Unlock() + return false + } +} + +func (s *channelFanoutShard) enqueueOverflow(channelID int64, pts int) bool { + s.mu.Lock() + accepted := s.addOverflowLocked(channelID, pts) + s.mu.Unlock() + if accepted { + s.signalOverflow() + } + return accepted +} + +func (s *channelFanoutShard) enqueueOverflowWait(ctx context.Context, channelID int64, pts int, stop <-chan struct{}) bool { + if ctx == nil { + ctx = context.Background() + } + select { + case <-ctx.Done(): + return false + case <-stop: + return false + default: + } + // Try and capture the current space generation under the same lock. Separating + // these operations creates a classic missed-wakeup window: a drain may free the + // mailbox after the failed try but before the waiter starts observing the signal. + s.mu.Lock() + if s.addOverflowLocked(channelID, pts) { + s.mu.Unlock() + s.signalOverflow() + return true + } + space := s.overflowSpace + s.overflowWaiters++ + s.mu.Unlock() + defer func() { + s.mu.Lock() + s.overflowWaiters-- + s.mu.Unlock() + }() + + // Same-channel overflow is the hot saturation path and only updates an existing map item. + // Distinct-channel saturation may wait here only from the dispatcher's fixed recovery sweep + // actor. RPC producers never call this method: they publish an O(1) global recovery generation + // when the bounded mailbox is full, so request goroutines cannot be exhausted by fan-out + // admission pressure. + for { + // Cardinality is full with distinct channels. Apply bounded-memory backpressure instead of + // allocating an unbounded recovery map or dropping the recovery watermark. + select { + case <-space: + case <-ctx.Done(): + return false + case <-stop: + return false + } + // Do not turn a release racing with cancellation into admission after the caller ended. + select { + case <-ctx.Done(): + return false + case <-stop: + return false + default: + } + // Retrying and subscribing to the next generation must also be atomic with + // respect to a release. Broadcast wakeups can be spurious for a particular + // waiter (another waiter may win the sole slot), so loop until accepted or stopped. + s.mu.Lock() + if s.addOverflowLocked(channelID, pts) { + s.mu.Unlock() + s.signalOverflow() + return true + } + space = s.overflowSpace + s.mu.Unlock() + } +} + +func (s *channelFanoutShard) addOverflowLocked(channelID int64, pts int) bool { + item, exists := s.overflow[channelID] + if !exists { + if len(s.overflow) >= s.overflowLimit { + return false + } + s.overflowOrder = append(s.overflowOrder, channelID) + // The first overflow fixes the FIFO barrier. Later same-channel losses only raise the + // durable pts watermark: moving the barrier on every merge lets a continuously full + // payload queue keep the recovery nudge one slot behind forever. UpdateChannelTooLong is + // an idempotent catch-up trigger, so it is safe for its newest pts to overtake payloads + // accepted after the first loss; those payloads become harmless duplicates after + // getChannelDifference converges the client. + item.barrier = s.nextSeq + } + if pts > item.pts { + item.pts = pts + } + s.overflow[channelID] = item + return true +} + +func (s *channelFanoutShard) markProcessed(seq uint64) { + s.mu.Lock() + if seq > s.processedSeq { + s.processedSeq = seq + } + s.mu.Unlock() +} + +// signalOverflowSpaceLocked announces a mailbox-cardinality decrease. Callers +// must hold s.mu. Close-and-replace provides broadcast generations without an +// unbounded waiter list, goroutine-per-waiter, or lossy fixed-capacity token queue. +func (s *channelFanoutShard) signalOverflowSpaceLocked() { + if s.overflowWaiters == 0 { + return + } + close(s.overflowSpace) + s.overflowSpace = make(chan struct{}) +} + +// popOverflow 仅供 mailbox/cardinality 单元测试直接释放一个 overflow；生产 drain 必须走 +// tryQueueOverflow，确保 nudgeJobs 真正接收成功前不删除水位。 +func (s *channelFanoutShard) popOverflow() (channelID int64, pts int, ok bool) { + s.mu.Lock() + defer s.mu.Unlock() + + for remaining := len(s.overflowOrder); remaining > 0; remaining-- { + channelID = s.overflowOrder[0] + s.overflowOrder = s.overflowOrder[1:] + item, exists := s.overflow[channelID] + if !exists { + continue + } + if item.barrier > s.processedSeq { + s.overflowOrder = append(s.overflowOrder, channelID) + continue + } + delete(s.overflow, channelID) + s.signalOverflowSpaceLocked() + return channelID, item.pts, true + } + return 0, 0, false +} + +// tryQueueOverflow 尝试把一个 barrier 已满足的 overflow 水位非阻塞提交给共享 nudge queue。 +// 只有 channel send 成功才从 mailbox 删除；queue 满时保留原 item（包含并发合并后的最高 pts） +// 和原 order 位置。整个操作在 shard.mu 下完成，因此不会出现“读到旧 pts 后删除新 pts”的竞态。 +func (s *channelFanoutShard) tryQueueOverflow(offer func(channelFanoutNudge) bool) (queued, blocked bool) { + s.mu.Lock() + defer s.mu.Unlock() + + for i := 0; i < len(s.overflowOrder); { + channelID := s.overflowOrder[i] + item, exists := s.overflow[channelID] + if !exists { + copy(s.overflowOrder[i:], s.overflowOrder[i+1:]) + s.overflowOrder = s.overflowOrder[:len(s.overflowOrder)-1] + continue + } + if item.barrier > s.processedSeq { + i++ + continue + } + if offer(channelFanoutNudge{channelID: channelID, pts: item.pts}) { + delete(s.overflow, channelID) + copy(s.overflowOrder[i:], s.overflowOrder[i+1:]) + s.overflowOrder = s.overflowOrder[:len(s.overflowOrder)-1] + s.signalOverflowSpaceLocked() + return true, false + } else { + // Shared nudge workers are saturated. Keep the exact watermark and retry from + // the shard's bounded timer; do not remove or advance the mailbox entry, and + // never park this payload worker on the nudge queue. + return false, true + } + } + return false, false +} + +func (s *channelFanoutShard) signalOverflow() { + select { + case s.overflowWake <- struct{}{}: + default: + } +} + +func (s *channelFanoutShard) signalEligibleOverflow() { + s.mu.Lock() + eligible := false + for _, channelID := range s.overflowOrder { + if item, ok := s.overflow[channelID]; ok && item.barrier <= s.processedSeq { + eligible = true + break + } + } + s.mu.Unlock() + if eligible { + s.signalOverflow() + } } // channelFanoutPrefetch 在 worker 解析出最终 recipient 集合后、逐 viewer build 之前调用一次， @@ -60,30 +342,72 @@ type channelFanoutPrefetch func(ctx context.Context, viewers []int64) // channelFanoutDispatcher 把频道 payload fan-out 移出发送者 RPC，按 channelID 分片串行处理。 type channelFanoutDispatcher struct { - r *Router - log *zap.Logger - shards []chan channelFanoutJob - started atomic.Bool + r *Router + log *zap.Logger + shards []*channelFanoutShard + started atomic.Bool + stopped atomic.Bool + stopCh chan struct{} + stopOnce sync.Once + enqueueMu sync.RWMutex + + budgetMu sync.Mutex + queuedJobs int + queuedBytes int64 + maxQueuedJobs int + maxQueuedBytes int64 + + nudgeJobs chan int64 + nudgeWorkers int + nudgeTimeout time.Duration + nudgeMu sync.Mutex + // nudgePending and nudgeJobs form one bounded coalescing mailbox. nudgeJobs contains only + // channel ids; the mutable map value always holds the highest pts observed before a worker + // takes that id. A hot channel therefore occupies one slot rather than filling the queue. + nudgePending map[int64]int + nudgeLimit int + + // recoveryGeneration is the terminal in-memory saturation fallback. It deliberately carries + // no channel id: the fixed recovery actor enumerates the online membership index and reloads + // each channel's durable max pts. Thus even when every key-bearing mailbox is full, publishing + // one constant-size generation cannot fail or block an RPC producer. + recoveryGeneration atomic.Uint64 + recoveryCompleted atomic.Uint64 + recoveryWake chan struct{} + // dropped 保留旧字段名供既有统计兼容；现在表示“真实 payload 因 queue full 被折叠为 + // nudge-only overflow watermark”的次数，不再表示恢复触发也被静默丢弃。 dropped atomic.Int64 } +type channelFanoutNudge struct { + channelID int64 + pts int +} + // enqueueChannelFanout 把一条 channel-payload-pts 的 fan-out 投入异步 dispatcher。 -// 从请求 ctx 抓取发起设备的业务 auth key + session_id 显式带入 job，使异步 worker 仍能 +// 从请求 ctx 抓取发起设备的 raw auth key + session_id 显式带入 job，使异步 worker 仍能 // 排除发起设备回显（请求 ctx 异步时已失效）。仅用于会推进客户端 channel PtsWaiter 的真实 // payload（新消息/编辑/删除/pin）；reaction/poll（viewer-only 零 pts）、participant/TTL/ // channel state（无 channel pts）、typing（transient）不走此路径（设计 §2.1/§5 分类）。 func (r *Router) enqueueChannelFanout(ctx context.Context, scope channelFanoutScope, originUserID, channelID int64, pts int, recipients []int64, build channelFanoutBuilder) { - r.enqueueChannelFanoutWithPrefetch(ctx, scope, originUserID, channelID, pts, recipients, nil, build) + r.enqueueChannelFanoutWithPrefetch(ctx, scope, originUserID, channelID, pts, recipients, 0, nil, build) } // enqueueChannelFanoutWithPrefetch 同 enqueueChannelFanout，但额外带一个跨 viewer 用户投影预热钩子 // （fan-out 模板化把每 recipient 的逐 viewer 投影折叠成一次 O(owner) 投影；见 prefetchChannelFanoutUsers）。 -func (r *Router) enqueueChannelFanoutWithPrefetch(ctx context.Context, scope channelFanoutScope, originUserID, channelID int64, pts int, recipients []int64, prefetch channelFanoutPrefetch, build channelFanoutBuilder) { +func (r *Router) enqueueChannelFanoutWithPrefetch(ctx context.Context, scope channelFanoutScope, originUserID, channelID int64, pts int, recipients []int64, retainedFloor int64, prefetch channelFanoutPrefetch, build channelFanoutBuilder) { if r.channelFanout == nil || build == nil { return } - originAuthKeyID, _ := AuthKeyIDFrom(ctx) + originAuthKeyID := rawAuthKeyIDForOrigin(ctx) originSessionID, _ := SessionIDFrom(ctx) + retainedBytes := int64(inboundRPCBytesFrom(ctx)) + int64(len(recipients))*8 + 4096 + if retainedFloor < channelFanoutMinRetainedBytes { + retainedFloor = channelFanoutMinRetainedBytes + } + if retainedBytes < retainedFloor { + retainedBytes = retainedFloor + } r.channelFanout.Enqueue(ctx, channelFanoutJob{ scope: scope, originUserID: originUserID, @@ -94,6 +418,7 @@ func (r *Router) enqueueChannelFanoutWithPrefetch(ctx context.Context, scope cha originSessionID: originSessionID, prefetch: prefetch, build: build, + retainedBytes: retainedBytes, }) } @@ -110,9 +435,22 @@ func newChannelFanoutDispatcher(r *Router, shards, buffer int) *channelFanoutDis if buffer <= 0 { buffer = defaultChannelFanoutBuffer } - d := &channelFanoutDispatcher{r: r, log: r.log.Named("channel-fanout"), shards: make([]chan channelFanoutJob, shards)} + d := &channelFanoutDispatcher{ + r: r, + log: r.log.Named("channel-fanout"), + shards: make([]*channelFanoutShard, shards), + stopCh: make(chan struct{}), + maxQueuedJobs: defaultChannelFanoutMaxQueuedJobs, + maxQueuedBytes: defaultChannelFanoutMaxQueuedBytes, + nudgeJobs: make(chan int64, defaultChannelNudgeQueue), + nudgeWorkers: defaultChannelNudgeWorkers, + nudgeTimeout: defaultChannelFanoutNudgeDeadline, + nudgePending: make(map[int64]int), + nudgeLimit: defaultChannelNudgeQueue, + recoveryWake: make(chan struct{}, 1), + } for i := range d.shards { - d.shards[i] = make(chan channelFanoutJob, buffer) + d.shards[i] = newChannelFanoutShard(buffer) } return d } @@ -124,22 +462,138 @@ func (d *channelFanoutDispatcher) Run(ctx context.Context) { return } var wg sync.WaitGroup - for i := range d.shards { + wg.Add(1) + go func() { + defer wg.Done() + <-ctx.Done() + d.enqueueMu.Lock() + d.stopped.Store(true) + d.stopOnce.Do(func() { close(d.stopCh) }) + d.enqueueMu.Unlock() + }() + wg.Add(1) + go func() { + defer wg.Done() + d.runRecoverySweeps(ctx) + }() + for range d.nudgeWorkers { wg.Add(1) - ch := d.shards[i] go func() { defer wg.Done() for { select { case <-ctx.Done(): return - case job := <-ch: + case channelID := <-d.nudgeJobs: + nudge, ok := d.takeNudge(channelID) + if !ok { + continue + } + timeout := d.nudgeTimeout + if timeout <= 0 { + timeout = defaultChannelFanoutNudgeDeadline + } + nudgeCtx, cancel := context.WithTimeout(ctx, timeout) + complete := d.r.runChannelFanoutOverflowNudge(nudgeCtx, nudge.channelID, nudge.pts) + cancel() + if !complete && ctx.Err() == nil { + // A deadline may leave only a prefix of online members nudged. Do not try to + // remember that recipient subset: request a durable max-pts sweep instead. + d.requestRecoverySweep("nudge deadline") + } + } + } + }() + } + for i := range d.shards { + wg.Add(1) + shard := d.shards[i] + go func() { + defer wg.Done() + var retryTimer *time.Timer + var retryC <-chan time.Time + retryDelay := channelFanoutNudgeRetryMin + stopRetryTimer := func() { + if retryTimer != nil { + retryTimer.Stop() + } + } + defer stopRetryTimer() + scheduleRetry := func() { + if retryC != nil { + return + } + if retryTimer == nil { + retryTimer = time.NewTimer(retryDelay) + } else { + retryTimer.Reset(retryDelay) + } + retryC = retryTimer.C + if retryDelay < channelFanoutNudgeRetryMax { + retryDelay *= 2 + if retryDelay > channelFanoutNudgeRetryMax { + retryDelay = channelFanoutNudgeRetryMax + } + } + } + drain := func() { + // While a retry is armed, payload completions and coalescing wakeups must not + // defeat backoff and spin on a full shared queue. + if retryC != nil { + return + } + queued, blocked := d.drainOneOverflow(shard) + if queued { + retryDelay = channelFanoutNudgeRetryMin + shard.signalEligibleOverflow() + return + } + if blocked { + scheduleRetry() + } + } + for { + select { + case <-ctx.Done(): + return + case job := <-shard.jobs: d.r.runChannelFanoutJob(ctx, job) + d.releaseQueuedJob(job) + shard.markProcessed(job.queueSeq) + // 每处理一条正常 FIFO payload，主动尝试 drain 一条已经越过 + // barrier 的 overflow。这样持续灌满正常队列的热点频道也不能 + // 永久饿死其它频道的恢复 nudge。 + drain() + case <-shard.overflowWake: + drain() + case <-retryC: + retryC = nil + drain() } } }() } wg.Wait() + // Workers may choose ctx.Done while jobs remain buffered. Release every reservation and + // drop closure references so tests/restarts do not retain the global budget after shutdown. + for _, shard := range d.shards { + for { + select { + case job := <-shard.jobs: + d.releaseQueuedJob(job) + default: + goto drained + } + } + drained: + shard.mu.Lock() + clear(shard.overflow) + shard.overflowOrder = nil + shard.mu.Unlock() + } + d.nudgeMu.Lock() + clear(d.nudgePending) + d.nudgeMu.Unlock() } func (d *channelFanoutDispatcher) shardIndex(channelID int64) int { @@ -152,8 +606,10 @@ func (d *channelFanoutDispatcher) shardIndex(channelID int64) int { } // Enqueue 投递一条 fan-out 任务。dispatcher 未启动时同步执行（用请求 ctx，保持旧行为）； -// 已启动时投入对应分片，满则丢弃 + 告警（该 channel 下一条消息的 pts 跳变会经 -// getChannelDifference 兜底）。 +// 已启动时投入对应分片。满时正常 payload 不阻塞请求路径，而是按 channel 合并为最高 pts +// 的 nudge-only overflow watermark，由同 shard worker 在更早的 FIFO payload 后公平 drain。 +// 若 overflow cardinality 也已满，只发布一个常量大小的全局 recovery generation；固定后台 +// actor 随后从 durable channel pts 重建全部在线 channel 的 nudge。RPC goroutine 永不等待 slot。 func (d *channelFanoutDispatcher) Enqueue(reqCtx context.Context, job channelFanoutJob) { if d == nil || job.build == nil { return @@ -162,14 +618,256 @@ func (d *channelFanoutDispatcher) Enqueue(reqCtx context.Context, job channelFan d.r.runChannelFanoutJob(reqCtx, job) return } - shard := d.shards[d.shardIndex(job.channelID)] - select { - case shard <- job: - default: - d.dropped.Add(1) - d.log.Warn("channel fanout queue full, dropped realtime push (recovered via next pts gap / getChannelDifference)", - zap.Int64("channel_id", job.channelID), zap.Int("pts", job.pts)) + d.enqueueMu.RLock() + if d.stopped.Load() { + d.enqueueMu.RUnlock() + return } + shard := d.shards[d.shardIndex(job.channelID)] + queued := false + if d.reserveQueuedJob(job) { + queued = shard.enqueue(job) + if queued { + d.enqueueMu.RUnlock() + return + } + d.releaseQueuedJob(job) + } + d.dropped.Add(1) + if job.scope != channelFanoutMembers || job.pts <= 0 { + // 当前所有 enqueue 入口均为 members + durable pts；若未来新增其它 scope，必须先 + // 定义其 overflow 恢复面，不能误把 viewer-only/no-pts 更新伪装成 channel nudge。 + d.log.Error("channel fanout queue full for non-coalescible job; overflow contract violated", + zap.Int64("channel_id", job.channelID), zap.Int("pts", job.pts), zap.Int("scope", int(job.scope))) + d.enqueueMu.RUnlock() + return + } + channelID, pts := job.channelID, job.pts + // The payload closure and recipient snapshot are no longer needed after normal queue + // admission failed. Make them unreachable before applying overflow-cardinality backpressure; + // otherwise blocked producers would retain unbudgeted request bodies while waiting for one of + // the fixed mailbox slots. Inbound RPC concurrency remains the producer-count bound. + job.recipients = nil + job.prefetch = nil + job.build = nil + if !shard.enqueueOverflow(channelID, pts) { + // Every key-bearing in-memory structure is bounded. Once the shard mailbox has no distinct + // channel slot, do not add another queue and do not park this RPC worker. A generation bit is + // enough because channels.pts/channel_update_events are already the durable truth: the fixed + // recovery actor can enumerate all online channel ids and reconstruct the highest watermark. + d.requestRecoverySweep("overflow cardinality full") + d.log.Warn("channel fanout overflow cardinality exhausted; scheduled durable max-pts recovery sweep", + zap.Int64("channel_id", channelID), zap.Int("pts", pts)) + d.enqueueMu.RUnlock() + return + } + d.log.Warn("channel fanout capacity exhausted, coalesced realtime payload into highest-pts overflow nudge", + zap.Int64("channel_id", channelID), zap.Int("pts", pts)) + d.enqueueMu.RUnlock() +} + +func (d *channelFanoutDispatcher) reserveQueuedJob(job channelFanoutJob) bool { + size := job.retainedBytes + if size < channelFanoutMinRetainedBytes { + size = channelFanoutMinRetainedBytes + } + d.budgetMu.Lock() + defer d.budgetMu.Unlock() + if d.queuedJobs >= d.maxQueuedJobs || size > d.maxQueuedBytes-d.queuedBytes { + return false + } + d.queuedJobs++ + d.queuedBytes += size + return true +} + +func (d *channelFanoutDispatcher) releaseQueuedJob(job channelFanoutJob) { + size := job.retainedBytes + if size < channelFanoutMinRetainedBytes { + size = channelFanoutMinRetainedBytes + } + d.budgetMu.Lock() + d.queuedJobs-- + d.queuedBytes -= size + if d.queuedJobs < 0 || d.queuedBytes < 0 { + panic("channel fanout queue budget underflow") + } + d.budgetMu.Unlock() +} + +func (d *channelFanoutDispatcher) queuedBudgetSnapshot() (jobs int, bytes int64) { + d.budgetMu.Lock() + defer d.budgetMu.Unlock() + return d.queuedJobs, d.queuedBytes +} + +func (d *channelFanoutDispatcher) drainOneOverflow(shard *channelFanoutShard) (queued, blocked bool) { + return shard.tryQueueOverflow(d.offerNudge) +} + +// offerNudge inserts one channel id into the bounded shared queue and stores its mutable highest +// pts in nudgePending. It never blocks. A same-channel update succeeds even when cardinality is +// full because it consumes no additional queue slot. +func (d *channelFanoutDispatcher) offerNudge(nudge channelFanoutNudge) bool { + if nudge.channelID == 0 || nudge.pts <= 0 { + return true + } + d.nudgeMu.Lock() + if current, exists := d.nudgePending[nudge.channelID]; exists { + if nudge.pts > current { + d.nudgePending[nudge.channelID] = nudge.pts + } + d.nudgeMu.Unlock() + return true + } + if len(d.nudgePending) >= d.nudgeLimit { + d.nudgeMu.Unlock() + return false + } + d.nudgePending[nudge.channelID] = nudge.pts + select { + case d.nudgeJobs <- nudge.channelID: + d.nudgeMu.Unlock() + return true + default: + // nudgeJobs has the same cardinality bound as nudgePending. This branch is reachable only + // while a test overrides one without the other or an invariant regresses; roll back rather + // than retain an unreachable map entry. + delete(d.nudgePending, nudge.channelID) + d.nudgeMu.Unlock() + return false + } +} + +func (d *channelFanoutDispatcher) takeNudge(channelID int64) (channelFanoutNudge, bool) { + d.nudgeMu.Lock() + pts, ok := d.nudgePending[channelID] + if ok { + delete(d.nudgePending, channelID) + } + d.nudgeMu.Unlock() + return channelFanoutNudge{channelID: channelID, pts: pts}, ok +} + +func (d *channelFanoutDispatcher) requestRecoverySweep(reason string) { + generation := d.recoveryGeneration.Add(1) + select { + case d.recoveryWake <- struct{}{}: + default: + } + d.log.Debug("channel fanout durable recovery sweep requested", + zap.Uint64("generation", generation), zap.String("reason", reason)) +} + +// runRecoverySweeps owns the only potentially waiting overflow admission path. Producers publish +// generations and return; this fixed actor reconstructs channel ids from the live membership index +// and watermarks from durable channels.pts. A generation is marked complete only after every page +// and every channel in that page has successfully entered its shard's barrier-preserving overflow +// mailbox. Errors retain the generation and retry with bounded backoff. +func (d *channelFanoutDispatcher) runRecoverySweeps(ctx context.Context) { + completed := d.recoveryCompleted.Load() + retryDelay := channelFanoutRecoveryRetryMin + var retryTimer *time.Timer + defer func() { + if retryTimer != nil { + retryTimer.Stop() + } + }() + for { + target := d.recoveryGeneration.Load() + if target <= completed { + select { + case <-ctx.Done(): + return + case <-d.recoveryWake: + continue + } + } + if err := d.sweepOnlineChannelRecovery(ctx); err != nil { + if ctx.Err() != nil { + return + } + d.log.Warn("channel fanout durable recovery sweep failed; retaining generation", + zap.Uint64("generation", target), zap.Duration("retry_in", retryDelay), zap.Error(err)) + if retryTimer == nil { + retryTimer = time.NewTimer(retryDelay) + } else { + retryTimer.Reset(retryDelay) + } + // recoveryGeneration already records every concurrent request. A wake may shorten the + // idle wait before a healthy sweep, but it must never bypass failure backoff: otherwise + // sustained saturation plus a persistent DB error retries at producer rate. + select { + case <-ctx.Done(): + return + case <-retryTimer.C: + } + if retryDelay < channelFanoutRecoveryRetryMax { + retryDelay *= 2 + if retryDelay > channelFanoutRecoveryRetryMax { + retryDelay = channelFanoutRecoveryRetryMax + } + } + continue + } + completed = target + d.recoveryCompleted.Store(completed) + retryDelay = channelFanoutRecoveryRetryMin + d.log.Info("channel fanout durable recovery sweep completed", zap.Uint64("generation", completed)) + // If a producer saturated after its channel had already been visited, generation is now + // greater than completed and the next loop immediately performs a fresh full pass. + } +} + +func (d *channelFanoutDispatcher) sweepOnlineChannelRecovery(ctx context.Context) error { + sessions, ok := d.r.deps.Sessions.(ChannelFanoutRecoverySessionProvider) + if !ok { + return fmt.Errorf("sessions dependency lacks online channel recovery enumeration") + } + channels, ok := d.r.deps.Channels.(ChannelFanoutRecoveryPtsProvider) + if !ok { + return fmt.Errorf("channels dependency lacks durable max pts lookup") + } + channelIDs := sessions.OnlineChannelIDsSnapshot() + for i, channelID := range channelIDs { + if channelID <= 0 || (i > 0 && channelID <= channelIDs[i-1]) { + return fmt.Errorf("online channel recovery snapshot is not strictly ascending: index=%d got=%d", i, channelID) + } + } + for start := 0; start < len(channelIDs); start += defaultChannelFanoutRecoverySweepPage { + end := start + defaultChannelFanoutRecoverySweepPage + if end > len(channelIDs) { + end = len(channelIDs) + } + page := channelIDs[start:end] + ptsByChannel, err := channels.MaxChannelPtsBatch(ctx, page) + if err != nil { + return fmt.Errorf("load durable max pts for online channel page [%d:%d]: %w", start, end, err) + } + for _, channelID := range page { + pts := ptsByChannel[channelID] + if pts > 0 { + shard := d.shards[d.shardIndex(channelID)] + if !shard.enqueueOverflowWait(ctx, channelID, pts, d.stopCh) { + if err := ctx.Err(); err != nil { + return err + } + return fmt.Errorf("dispatcher stopped while admitting recovery for channel %d", channelID) + } + } + } + } + return nil +} + +// runChannelFanoutOverflowNudge 是 queue-full 的 nudge-only 降级路径。不能复用原 job 的 +// origin exclude：同一 channel 水位可能合并多个不同发起 session；向全部在线成员发最高 pts +// nudge 是保守且幂等的，已追上 pts 的 TDesktop 会直接忽略。 +func (r *Router) runChannelFanoutOverflowNudge(ctx context.Context, channelID int64, pts int) bool { + if r.deps.Sessions == nil || channelID == 0 || pts <= 0 { + return true + } + return r.nudgeBeyondCapChannelMembers(ctx, channelID, pts, nil) } // runChannelFanoutJob 执行一条 fan-out：与同步 pushChannelUpdatesWithScope 等价，区别是 @@ -179,7 +877,7 @@ func (r *Router) runChannelFanoutJob(ctx context.Context, job channelFanoutJob) if r.deps.Sessions == nil || job.build == nil { return } - pushCtx := WithSessionID(WithAuthKeyID(ctx, job.originAuthKeyID), job.originSessionID) + pushCtx := WithSessionID(WithRawAuthKeyID(ctx, job.originAuthKeyID), job.originSessionID) recipients := r.channelFanoutRecipients(ctx, job.scope, job.channelID, job.recipients) // 预热跨 viewer 用户投影（fan-out 模板化）：在逐 viewer build 之前一次性算好每 recipient 的 // 投影并预热共享 cache，使 build 只命中缓存、不再 O(viewer) 逐个 ForViewer。覆盖 recipients + @@ -278,6 +976,7 @@ func (r *Router) enqueueChannelMessageFanout(ctx context.Context, originUserID i ownerIDs := channelMessageFanoutOwnerIDs(res, extraUserIDs) skip := skipDeliverySet(res.SkipDeliveryUserIDs) r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, res.Channel.ID, res.Event.Pts, res.Recipients, + 0, func(bgCtx context.Context, viewers []int64) { r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs) }, @@ -342,6 +1041,7 @@ func (r *Router) enqueueChannelEditMessageFanout(ctx context.Context, originUser ownerIDs := channelEditMessageFanoutOwnerIDs(res) nudgePts := max(res.Event.Pts, res.ServiceEvent.Pts) r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, res.Channel.ID, nudgePts, res.Recipients, + 0, func(bgCtx context.Context, viewers []int64) { r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs) }, @@ -358,6 +1058,7 @@ func (r *Router) enqueueChannelMessagesFanout(ctx context.Context, originUserID, fanoutCache := newViewerPeerCache(r) ownerIDs := channelMessagesFanoutOwnerIDs(results, extraUserIDs) r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, channelID, pts, recipients, + int64(len(results))*(64<<10), func(bgCtx context.Context, viewers []int64) { r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs) }, @@ -385,28 +1086,37 @@ func (r *Router) channelNudgeMaxTargets() int { // getChannelDifference（设计 §10.3）。走 pushUserUpdates（best-effort、未就绪入 pending、非 // transient），符合设计 §决策4 的 nudge 投递可靠性要求。SessionManager 未实现 ChannelNudgeProvider // 时（测试/未装配）静默跳过，不影响完整 payload 投递。 -func (r *Router) nudgeBeyondCapChannelMembers(ctx context.Context, channelID int64, pts int, delivered map[int64]struct{}) { +func (r *Router) nudgeBeyondCapChannelMembers(ctx context.Context, channelID int64, pts int, delivered map[int64]struct{}) bool { provider, ok := r.deps.Sessions.(ChannelNudgeProvider) if !ok || channelID == 0 || pts <= 0 { - return + return true } targets := provider.OnlineChannelMemberUserIDsExcluding(channelID, delivered, r.channelNudgeMaxTargets()) if len(targets) == 0 { - return + return true } date := int(r.clock.Now().Unix()) + tooLong := &tg.UpdateChannelTooLong{ChannelID: channelID} + tooLong.SetPts(pts) + updates := &tg.Updates{ + Updates: []tg.UpdateClass{tooLong}, + Users: []tg.UserClass{}, + Chats: []tg.ChatClass{}, + Date: date, + Seq: 0, + } for _, userID := range targets { + select { + case <-ctx.Done(): + return false + default: + } if userID == 0 { continue } - tooLong := &tg.UpdateChannelTooLong{ChannelID: channelID} - tooLong.SetPts(pts) - r.pushUserUpdates(ctx, userID, &tg.Updates{ - Updates: []tg.UpdateClass{tooLong}, - Users: []tg.UserClass{}, - Chats: []tg.ChatClass{}, - Date: date, - Seq: 0, - }) + // The nudge is viewer-independent and immutable. Reuse the TL object across + // recipients; SessionManager encodes before enqueue and never mutates it. + r.pushUserUpdates(ctx, userID, updates) } + return ctx.Err() == nil } diff --git a/internal/rpc/channel_fanout_dispatcher_test.go b/internal/rpc/channel_fanout_dispatcher_test.go index 1cc0cd5d..ec4f8b1a 100644 --- a/internal/rpc/channel_fanout_dispatcher_test.go +++ b/internal/rpc/channel_fanout_dispatcher_test.go @@ -2,6 +2,8 @@ package rpc import ( "context" + "fmt" + "runtime" "sync" "testing" "time" @@ -41,6 +43,150 @@ func fanoutHasID(ids []int64, want int64) bool { return false } +type recoveryFanoutSessions struct { + SessionBinder + + onlineChannels []int64 + pushStarted chan struct{} + pushRelease <-chan struct{} + startOnce sync.Once + + mu sync.Mutex + nudges map[int64][]int + pushErr int +} + +func newRecoveryFanoutSessions(onlineChannels []int64, release <-chan struct{}) *recoveryFanoutSessions { + return &recoveryFanoutSessions{ + onlineChannels: append([]int64(nil), onlineChannels...), + pushStarted: make(chan struct{}), + pushRelease: release, + nudges: make(map[int64][]int), + } +} + +func (s *recoveryFanoutSessions) PushToUserExceptSession(ctx context.Context, _ int64, _ int64, _ proto.MessageType, msg bin.Encoder) (int, error) { + s.startOnce.Do(func() { close(s.pushStarted) }) + if s.pushRelease != nil { + select { + case <-ctx.Done(): + s.mu.Lock() + s.pushErr++ + s.mu.Unlock() + return 0, ctx.Err() + case <-s.pushRelease: + } + } + updates, ok := msg.(*tg.Updates) + if !ok || len(updates.Updates) != 1 { + return 1, nil + } + nudge, ok := updates.Updates[0].(*tg.UpdateChannelTooLong) + if !ok { + return 1, nil + } + pts, _ := nudge.GetPts() + s.mu.Lock() + s.nudges[nudge.ChannelID] = append(s.nudges[nudge.ChannelID], pts) + s.mu.Unlock() + return 1, nil +} + +func (s *recoveryFanoutSessions) OnlineChannelMemberUserIDsExcluding(_ int64, _ map[int64]struct{}, _ int) []int64 { + return []int64{42} +} + +func (s *recoveryFanoutSessions) OnlineChannelIDsAfter(afterChannelID int64, limit int) []int64 { + out := make([]int64, 0, limit) + for _, channelID := range s.onlineChannels { + if channelID <= afterChannelID { + continue + } + out = append(out, channelID) + if len(out) == limit { + break + } + } + return out +} + +func (s *recoveryFanoutSessions) OnlineChannelIDsSnapshot() []int64 { + return append([]int64(nil), s.onlineChannels...) +} + +func (s *recoveryFanoutSessions) nudgePts(channelID int64) []int { + s.mu.Lock() + defer s.mu.Unlock() + return append([]int(nil), s.nudges[channelID]...) +} + +func (s *recoveryFanoutSessions) pushErrors() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.pushErr +} + +type recoveryFanoutChannels struct { + ChannelsService + + mu sync.Mutex + pts map[int64]int + calls int + failCalls int + firstCalled chan struct{} + firstOnce sync.Once + release <-chan struct{} +} + +func (s *recoveryFanoutChannels) MaxChannelPts(ctx context.Context, channelID int64) (int, error) { + pts, err := s.MaxChannelPtsBatch(ctx, []int64{channelID}) + return pts[channelID], err +} + +func (s *recoveryFanoutChannels) MaxChannelPtsBatch(ctx context.Context, channelIDs []int64) (map[int64]int, error) { + s.mu.Lock() + s.calls++ + call := s.calls + if call <= s.failCalls { + s.mu.Unlock() + return nil, fmt.Errorf("injected max pts failure %d", call) + } + pts := make(map[int64]int, len(channelIDs)) + for _, channelID := range channelIDs { + if value, ok := s.pts[channelID]; ok { + pts[channelID] = value + } + } + release := s.release + s.mu.Unlock() + if s.firstCalled != nil { + s.firstOnce.Do(func() { close(s.firstCalled) }) + } + if release != nil { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-release: + } + } + return pts, nil +} + +func (s *recoveryFanoutChannels) setPts(channelID int64, pts int) { + s.mu.Lock() + if s.pts == nil { + s.pts = make(map[int64]int) + } + s.pts[channelID] = pts + s.mu.Unlock() +} + +func (s *recoveryFanoutChannels) callCount() int { + s.mu.Lock() + defer s.mu.Unlock() + return s.calls +} + // TestChannelFanoutDispatcherSyncFallback：dispatcher 未启动时 Enqueue 同步执行—— // 保持测试/未装配场景行为不变，recipients 立即被推送、发起 session 作为 exclude 透传。 // deps.Channels=nil 时 channelFanoutRecipients 直接返回 explicit recipients。 @@ -232,6 +378,69 @@ type nudgeSessions struct { byUser map[int64]bin.Encoder } +// overflowNudgeSessions 为 queue-full 回归按 channel 提供不同在线成员，并只记录 +// UpdateChannelTooLong。正常 FIFO payload 仍委托 captureSessions 记录，但不会污染 nudge 断言。 +type overflowNudgeSessions struct { + *captureSessions + onlineByChannel map[int64][]int64 + + mu sync.Mutex + nudges map[int64][]int + order []string +} + +func newOverflowNudgeSessions(onlineByChannel map[int64][]int64) *overflowNudgeSessions { + return &overflowNudgeSessions{ + captureSessions: &captureSessions{}, + onlineByChannel: onlineByChannel, + nudges: make(map[int64][]int), + } +} + +func (s *overflowNudgeSessions) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, typ proto.MessageType, msg bin.Encoder) (int, error) { + if updates, ok := msg.(*tg.Updates); ok && len(updates.Updates) == 1 { + if nudge, ok := updates.Updates[0].(*tg.UpdateChannelTooLong); ok { + pts, _ := nudge.GetPts() + s.mu.Lock() + s.nudges[nudge.ChannelID] = append(s.nudges[nudge.ChannelID], pts) + s.order = append(s.order, fmt.Sprintf("nudge:%d:%d", nudge.ChannelID, pts)) + s.mu.Unlock() + } + } + return s.captureSessions.PushToUserExceptSession(ctx, userID, excludeSessionID, typ, msg) +} + +func (s *overflowNudgeSessions) OnlineChannelMemberUserIDsExcluding(channelID int64, exclude map[int64]struct{}, limit int) []int64 { + online := s.onlineByChannel[channelID] + out := make([]int64, 0, len(online)) + for _, id := range online { + if _, ok := exclude[id]; ok { + continue + } + out = append(out, id) + if limit > 0 && len(out) >= limit { + break + } + } + return out +} + +func (s *overflowNudgeSessions) recordOrder(event string) { + s.mu.Lock() + s.order = append(s.order, event) + s.mu.Unlock() +} + +func (s *overflowNudgeSessions) snapshot() (map[int64][]int, []string) { + s.mu.Lock() + defer s.mu.Unlock() + nudges := make(map[int64][]int, len(s.nudges)) + for channelID, pts := range s.nudges { + nudges[channelID] = append([]int(nil), pts...) + } + return nudges, append([]string(nil), s.order...) +} + func newNudgeSessions(online []int64) *nudgeSessions { return &nudgeSessions{captureSessions: &captureSessions{}, online: online, byUser: map[int64]bin.Encoder{}} } @@ -317,3 +526,939 @@ func TestChannelEditMessageFanoutNudgePtsUsesMaxContainer(t *testing.T) { t.Fatalf("nudge pts=%d ok=%v, want 11 (max(Event=0, Service=11))", p, ok) } } + +// TestChannelFanoutDispatcherOverflowCoalescesHighestPtsAndDrains 验证 queue full 不再静默 +// 丢掉恢复触发：正常 payload 仍按 FIFO 执行；同 channel 多次 overflow 只发最高 pts nudge； +// 另一个落在同 shard 的 channel 也最终得到 nudge，且不需要后续再 Enqueue 才唤醒 worker。 +func TestChannelFanoutDispatcherOverflowCoalescesHighestPtsAndDrains(t *testing.T) { + const ( + hotChannel = int64(1001) + otherChannel = int64(2001) + hotUser = int64(10001) + otherUser = int64(20001) + ) + sessions := newOverflowNudgeSessions(map[int64][]int64{ + hotChannel: {hotUser}, + otherChannel: {otherUser}, + }) + r := New(Config{}, Deps{Sessions: sessions}, zaptest.NewLogger(t), clock.System) + // 单 shard + 单 buffer 让测试确定地产生：1 条执行中、1 条正常 FIFO pending、其余 overflow。 + r.channelFanout = newChannelFanoutDispatcher(r, 1, 1) + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + r.RunChannelFanout(ctx) + close(done) + }() + for i := 0; i < 200 && !r.channelFanout.started.Load(); i++ { + time.Sleep(time.Millisecond) + } + if !r.channelFanout.started.Load() { + cancel() + <-done + t.Fatal("dispatcher did not start") + } + + firstStarted := make(chan struct{}) + releaseFirst := make(chan struct{}) + overflowBuildCalled := make(chan struct{}, 1) + job := func(channelID int64, pts int, recipients []int64, build channelFanoutBuilder) channelFanoutJob { + return channelFanoutJob{ + scope: channelFanoutMembers, + channelID: channelID, + pts: pts, + recipients: recipients, + build: build, + } + } + payload := func(channelID int64) *tg.Updates { + return &tg.Updates{ + Updates: []tg.UpdateClass{&tg.UpdateChannel{ChannelID: channelID}}, + Users: []tg.UserClass{}, + Chats: []tg.ChatClass{}, + Date: 1, + } + } + + r.channelFanout.Enqueue(context.Background(), job(hotChannel, 1, []int64{hotUser}, func(context.Context, int64) *tg.Updates { + close(firstStarted) + <-releaseFirst + sessions.recordOrder("payload:1001:1") + return payload(hotChannel) + })) + select { + case <-firstStarted: + case <-time.After(2 * time.Second): + cancel() + <-done + t.Fatal("first FIFO payload did not start") + } + + // 这条占满唯一正常 buffer；overflow barrier 必须等它完成后才允许发 nudge。 + r.channelFanout.Enqueue(context.Background(), job(hotChannel, 2, []int64{hotUser}, func(context.Context, int64) *tg.Updates { + sessions.recordOrder("payload:1001:2") + return payload(hotChannel) + })) + overflowBuild := func(context.Context, int64) *tg.Updates { + select { + case overflowBuildCalled <- struct{}{}: + default: + } + return nil + } + // 热点 channel 连续灌入只占一个 overflow mailbox 项，最终仅保留 pts=20。 + for pts := 3; pts <= 20; pts++ { + r.channelFanout.Enqueue(context.Background(), job(hotChannel, pts, []int64{hotUser}, overflowBuild)) + } + // 同 shard 的其它 channel 也必须最终 drain，不能被热点 channel 永久饿死。 + r.channelFanout.Enqueue(context.Background(), job(otherChannel, 7, []int64{otherUser}, overflowBuild)) + if got, want := r.channelFanout.dropped.Load(), int64(19); got != want { + cancel() + <-done + t.Fatalf("coalesced overflow jobs = %d, want %d", got, want) + } + + close(releaseFirst) + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + nudges, _ := sessions.snapshot() + if len(nudges[hotChannel]) == 1 && len(nudges[otherChannel]) == 1 { + break + } + time.Sleep(2 * time.Millisecond) + } + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("dispatcher did not stop") + } + + nudges, order := sessions.snapshot() + if got := nudges[hotChannel]; len(got) != 1 || got[0] != 20 { + t.Fatalf("hot channel nudges = %v, want one highest-pts nudge [20]", got) + } + if got := nudges[otherChannel]; len(got) != 1 || got[0] != 7 { + t.Fatalf("other channel nudges = %v, want [7] (must not be starved by hot channel)", got) + } + select { + case <-overflowBuildCalled: + t.Fatal("overflow job build was called; queue-full path must be nudge-only") + default: + } + if len(order) != 4 { + t.Fatalf("delivery order = %v, want two payloads + two nudges", order) + } + if order[0] != "payload:1001:1" || order[1] != "payload:1001:2" { + t.Fatalf("delivery order = %v, normal payload FIFO/barrier violated", order) + } + wantNudges := map[string]bool{"nudge:1001:20": true, "nudge:2001:7": true} + if !wantNudges[order[2]] || !wantNudges[order[3]] || order[2] == order[3] { + t.Fatalf("delivery order = %v, want both nudges after FIFO barrier (nudge pool may reorder channels)", order) + } +} + +func TestChannelFanoutQueueBudgetIsBoundedAndPreciselyReleased(t *testing.T) { + r := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.maxQueuedJobs = 1 + d.maxQueuedBytes = channelFanoutMinRetainedBytes + job := fanoutTestJob(nil, 0, 0, nil) + job.retainedBytes = channelFanoutMinRetainedBytes + if !d.reserveQueuedJob(job) { + t.Fatal("first job reservation rejected") + } + if d.reserveQueuedJob(job) { + t.Fatal("second job reservation exceeded global count/byte budget") + } + if jobs, bytes := d.queuedBudgetSnapshot(); jobs != 1 || bytes != channelFanoutMinRetainedBytes { + t.Fatalf("budget = %d/%d, want 1/%d", jobs, bytes, channelFanoutMinRetainedBytes) + } + d.releaseQueuedJob(job) + if jobs, bytes := d.queuedBudgetSnapshot(); jobs != 0 || bytes != 0 { + t.Fatalf("released budget = %d/%d, want 0/0", jobs, bytes) + } +} + +func TestChannelFanoutOverflowCardinalityBackpressuresUntilSpace(t *testing.T) { + shard := newChannelFanoutShard(1) + shard.overflowLimit = 1 + if !shard.enqueueOverflow(1001, 5) { + t.Fatal("first overflow watermark rejected") + } + stop := make(chan struct{}) + accepted := make(chan bool, 1) + go func() { accepted <- shard.enqueueOverflowWait(context.Background(), 2001, 7, stop) }() + select { + case <-accepted: + t.Fatal("second unique channel bypassed overflow cardinality bound") + case <-time.After(20 * time.Millisecond): + } + channelID, pts, ok := shard.popOverflow() + if !ok || channelID != 1001 || pts != 5 { + t.Fatalf("first pop = %d/%d/%v, want 1001/5/true", channelID, pts, ok) + } + select { + case ok := <-accepted: + if !ok { + t.Fatal("waiting overflow rejected after space became available") + } + case <-time.After(time.Second): + t.Fatal("waiting overflow did not resume after space became available") + } + channelID, pts, ok = shard.popOverflow() + if !ok || channelID != 2001 || pts != 7 { + t.Fatalf("second pop = %d/%d/%v, want 2001/7/true", channelID, pts, ok) + } +} + +func TestChannelFanoutOverflowSpaceBroadcastWakesAllWaitersAfterBatchRelease(t *testing.T) { + const waiters = 8 + shard := newChannelFanoutShard(1) + shard.overflowLimit = waiters + for i := range waiters { + if !shard.enqueueOverflow(int64(1000+i), i+1) { + t.Fatalf("initial overflow watermark %d rejected", i) + } + } + + results := make(chan bool, waiters) + for i := range waiters { + go func(i int) { + results <- shard.enqueueOverflowWait(context.Background(), int64(2000+i), 100+i, make(chan struct{})) + }(i) + } + + // Wait until every goroutine has atomically observed the same full-mailbox + // generation. This makes the regression deterministic: a one-token space + // notification can admit at most one waiter after the batch release below. + waitDeadline := time.Now().Add(50 * time.Millisecond) + for { + shard.mu.Lock() + waiting := shard.overflowWaiters + shard.mu.Unlock() + if waiting == waiters { + break + } + if time.Now().After(waitDeadline) { + t.Fatalf("overflow waiters = %d, want %d before release", waiting, waiters) + } + time.Sleep(time.Millisecond) + } + + // Model a drain that releases several cardinality slots before any waiter can + // run. One generation broadcast is sufficient: waiters serialize under mu and + // consume the eight real slots; no notification count is used as capacity. + shard.mu.Lock() + for channelID := range shard.overflow { + delete(shard.overflow, channelID) + } + shard.overflowOrder = shard.overflowOrder[:0] + shard.signalOverflowSpaceLocked() + shard.mu.Unlock() + + for i := range waiters { + select { + case accepted := <-results: + if !accepted { + t.Fatalf("waiter %d timed out after batch space release", i) + } + case <-time.After(time.Second): + t.Fatalf("waiter %d remained blocked after batch space release", i) + } + } + shard.mu.Lock() + mailboxLen := len(shard.overflow) + remainingWaiters := shard.overflowWaiters + shard.mu.Unlock() + if mailboxLen != waiters || remainingWaiters != 0 { + t.Fatalf("post-release mailbox/waiters = %d/%d, want %d/0", mailboxLen, remainingWaiters, waiters) + } +} + +func TestChannelFanoutSameKeyOverflowKeepsFirstBarrierUnderContinuousPayload(t *testing.T) { + const channelID = int64(1001) + shard := newChannelFanoutShard(1) + firstJob := channelFanoutJob{channelID: channelID, pts: 1} + if !shard.enqueue(firstJob) { + t.Fatal("first payload enqueue rejected") + } + if !shard.enqueueOverflow(channelID, 2) { + t.Fatal("first overflow watermark rejected") + } + + // Model a continuously saturated producer: as soon as the worker takes one payload, another + // payload occupies the slot and a newer loss merges into the same overflow key. The recovery + // nudge must be eligible after the original barrier, without waiting for the producer to stop. + processed := <-shard.jobs + shard.markProcessed(processed.queueSeq) + if !shard.enqueue(channelFanoutJob{channelID: channelID, pts: 3}) { + t.Fatal("replacement payload enqueue rejected") + } + if !shard.enqueueOverflow(channelID, 4) { + t.Fatal("same-key overflow merge rejected") + } + + shard.mu.Lock() + item := shard.overflow[channelID] + nextSeq := shard.nextSeq + shard.mu.Unlock() + if item.barrier != processed.queueSeq || nextSeq <= item.barrier { + t.Fatalf("overflow barrier/next = %d/%d, want first barrier %d while a newer payload remains queued", item.barrier, nextSeq, processed.queueSeq) + } + var got channelFanoutNudge + queued, blocked := shard.tryQueueOverflow(func(nudge channelFanoutNudge) bool { + got = nudge + return true + }) + if !queued || blocked || got.channelID != channelID || got.pts != 4 { + t.Fatalf("continuous-load drain = queued:%v blocked:%v nudge:%+v, want highest pts 4", queued, blocked, got) + } + if len(shard.jobs) != 1 { + t.Fatalf("replacement payload queue length = %d, want producer still active", len(shard.jobs)) + } +} + +func TestChannelFanoutRecoveryFailureBackoffIgnoresContinuousWake(t *testing.T) { + sessions := newRecoveryFanoutSessions([]int64{10}, nil) + channels := &recoveryFanoutChannels{pts: map[int64]int{10: 7}, failCalls: 1 << 30} + r := New(Config{}, Deps{Sessions: sessions, Channels: channels}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.nudgeWorkers = 0 + r.channelFanout = d + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { d.Run(ctx); close(done) }() + for !d.started.Load() { + time.Sleep(time.Millisecond) + } + + deadline := time.Now().Add(85 * time.Millisecond) + for time.Now().Before(deadline) { + d.requestRecoverySweep("continuous saturation during injected DB failure") + time.Sleep(time.Millisecond) + } + calls := channels.callCount() + cancel() + <-done + // With 10ms -> 20ms -> 40ms backoff this window permits about four calls. Keep a generous + // scheduler margin; the old wake-bypasses-timer loop makes tens of calls in the same window. + if calls < 2 || calls > 7 { + t.Fatalf("MaxChannelPtsBatch calls under continuous wake = %d, want bounded backoff in [2,7]", calls) + } +} + +func TestChannelFanoutShutdownWaitsForInFlightOverflowAdmission(t *testing.T) { + r := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.maxQueuedJobs = 0 // force Enqueue directly into overflow admission + r.channelFanout = d + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { d.Run(ctx); close(done) }() + for !d.started.Load() { + time.Sleep(time.Millisecond) + } + + shard := d.shards[0] + shard.mu.Lock() // hold admission so the enqueue lifetime is observable + enqueueDone := make(chan struct{}) + go func() { + d.Enqueue(context.Background(), channelFanoutJob{ + scope: channelFanoutMembers, channelID: 1001, pts: 5, + build: func(context.Context, int64) *tg.Updates { return nil }, + }) + close(enqueueDone) + }() + observedReader := false + waitUntil := time.Now().Add(time.Second) + for time.Now().Before(waitUntil) { + if !d.enqueueMu.TryLock() { + observedReader = true + break + } + d.enqueueMu.Unlock() + time.Sleep(time.Millisecond) + } + if !observedReader { + shard.mu.Unlock() + cancel() + <-done + t.Fatal("Enqueue did not retain shutdown gate while waiting for overflow admission") + } + + cancel() + time.Sleep(20 * time.Millisecond) + if d.stopped.Load() { + shard.mu.Unlock() + <-enqueueDone + <-done + t.Fatal("shutdown crossed an in-flight overflow admission") + } + shard.mu.Unlock() + <-enqueueDone + <-done + shard.mu.Lock() + remaining := len(shard.overflow) + shard.mu.Unlock() + if remaining != 0 { + t.Fatalf("overflow entries after shutdown = %d, want cleanup after all admissions finish", remaining) + } +} + +func TestChannelFanoutNudgeQueueFullRetainsHighestPtsUntilRetrySucceeds(t *testing.T) { + const channelID = int64(1001) + r := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.nudgeWorkers = 0 + d.nudgeJobs = make(chan int64, 1) + d.nudgeLimit = 1 + if !d.offerNudge(channelFanoutNudge{channelID: 9999, pts: 1}) { + t.Fatal("failed to prefill nudge mailbox") + } + r.channelFanout = d + + shard := d.shards[0] + if !shard.enqueueOverflow(channelID, 5) { + t.Fatal("initial overflow watermark rejected") + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + d.Run(ctx) + close(done) + }() + for i := 0; i < 200 && !d.started.Load(); i++ { + time.Sleep(time.Millisecond) + } + if !d.started.Load() { + cancel() + <-done + t.Fatal("dispatcher did not start") + } + + // Let the shard observe the full shared queue, then merge newer/lower pts while retries + // remain backpressured. The failed queue attempt must not remove the mailbox entry. + time.Sleep(10 * time.Millisecond) + if !shard.enqueueOverflow(channelID, 9) || !shard.enqueueOverflow(channelID, 7) { + cancel() + <-done + t.Fatal("same-channel overflow merge was rejected") + } + time.Sleep(10 * time.Millisecond) + shard.mu.Lock() + item, exists := shard.overflow[channelID] + mailboxLen := len(shard.overflow) + shard.mu.Unlock() + if !exists || mailboxLen != 1 || item.pts != 9 { + cancel() + <-done + t.Fatalf("full nudge queue mailbox = exists:%v len:%d item:%+v, want one retained pts=9", exists, mailboxLen, item) + } + + // Free one shared slot. The shard-owned bounded retry timer must submit the retained + // highest watermark without requiring another payload or Enqueue call. + blockedID := <-d.nudgeJobs + if _, ok := d.takeNudge(blockedID); !ok { + t.Fatal("prefilled nudge mailbox lost its pts") + } + select { + case gotID := <-d.nudgeJobs: + got, ok := d.takeNudge(gotID) + if !ok { + cancel() + <-done + t.Fatal("retried nudge queue id had no coalesced pts") + } + if got.channelID != channelID || got.pts != 9 { + cancel() + <-done + t.Fatalf("retried nudge = %+v, want channel=%d pts=9", got, channelID) + } + case <-time.After(time.Second): + cancel() + <-done + t.Fatal("retained overflow was not retried after nudge queue space became available") + } + shard.mu.Lock() + _, exists = shard.overflow[channelID] + shard.mu.Unlock() + if exists { + cancel() + <-done + t.Fatal("overflow watermark remained after successful nudge queue submission") + } + + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("dispatcher did not stop with an armed/recent nudge retry") + } +} + +func TestChannelFanoutNudgeBackpressureDoesNotBlockPayloadShards(t *testing.T) { + const ( + shardCount = 4 + jobsPerShard = 64 + ) + sessions := &captureSessions{} + r := New(Config{}, Deps{Sessions: sessions}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, shardCount, jobsPerShard+1) + d.nudgeWorkers = 0 + d.nudgeJobs = make(chan int64, 1) + d.nudgeLimit = 1 + if !d.offerNudge(channelFanoutNudge{channelID: 9999, pts: 1}) { + t.Fatal("failed to prefill nudge mailbox") + } + r.channelFanout = d + + channelIDs := make([]int64, shardCount) + for shardIndex := range shardCount { + channelID := int64(1000 + shardIndex) + for d.shardIndex(channelID) != shardIndex { + channelID++ + } + channelIDs[shardIndex] = channelID + if !d.shards[shardIndex].enqueueOverflow(channelID, 10+shardIndex) { + t.Fatalf("shard %d overflow watermark rejected", shardIndex) + } + } + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + d.Run(ctx) + close(done) + }() + for i := 0; i < 200 && !d.started.Load(); i++ { + time.Sleep(time.Millisecond) + } + if !d.started.Load() { + cancel() + <-done + t.Fatal("dispatcher did not start") + } + + processed := make(chan struct{}, shardCount*jobsPerShard) + for shardIndex, channelID := range channelIDs { + for pts := 1; pts <= jobsPerShard; pts++ { + d.Enqueue(context.Background(), channelFanoutJob{ + scope: channelFanoutMembers, + channelID: channelID, + pts: pts, + recipients: []int64{int64(2000 + shardIndex)}, + build: func(context.Context, int64) *tg.Updates { + processed <- struct{}{} + return &tg.Updates{Updates: []tg.UpdateClass{}, Users: []tg.UserClass{}, Chats: []tg.ChatClass{}, Date: 1} + }, + }) + } + } + for completed := 0; completed < shardCount*jobsPerShard; completed++ { + select { + case <-processed: + case <-time.After(2 * time.Second): + cancel() + <-done + t.Fatalf("payload shards stalled at %d/%d while nudge queue was full", completed, shardCount*jobsPerShard) + } + } + for shardIndex, shard := range d.shards { + shard.mu.Lock() + item, exists := shard.overflow[channelIDs[shardIndex]] + shard.mu.Unlock() + if !exists || item.pts != 10+shardIndex { + cancel() + <-done + t.Fatalf("shard %d lost overflow under nudge backpressure: exists=%v item=%+v", shardIndex, exists, item) + } + } + + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("dispatcher stop blocked behind a full nudge queue") + } + if jobs, bytes := d.queuedBudgetSnapshot(); jobs != 0 || bytes != 0 { + t.Fatalf("shutdown queue budget = %d/%d, want 0/0", jobs, bytes) + } +} + +func TestChannelFanoutAllMemoryLayersFullReturnsRPCAndDurableSweepRecovers(t *testing.T) { + const ( + fullChannel = int64(1001) + lostChannel = int64(3001) + lostMaxPts = 123 + ) + nudgeRelease := make(chan struct{}) + sessions := newRecoveryFanoutSessions([]int64{lostChannel}, nudgeRelease) + channels := &recoveryFanoutChannels{pts: map[int64]int{lostChannel: lostMaxPts}} + r := New(Config{}, Deps{Sessions: sessions, Channels: channels}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.shards[0].overflowLimit = 1 + d.nudgeWorkers = 1 + d.nudgeJobs = make(chan int64, 1) + d.nudgeLimit = 1 + r.channelFanout = d + + // Occupy the sole nudge worker, then the sole queued nudge slot. + if !d.offerNudge(channelFanoutNudge{channelID: 9000, pts: 1}) { + t.Fatal("initial nudge rejected") + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + d.Run(ctx) + close(done) + }() + select { + case <-sessions.pushStarted: + case <-time.After(time.Second): + cancel() + <-done + t.Fatal("nudge worker did not enter blocking push") + } + if !d.offerNudge(channelFanoutNudge{channelID: 9001, pts: 2}) { + cancel() + <-done + t.Fatal("queued nudge slot was not available") + } + + // Occupy the payload worker and its one buffered slot. + payloadStarted := make(chan struct{}) + payloadRelease := make(chan struct{}) + job := func(channelID int64, pts int, build channelFanoutBuilder) channelFanoutJob { + return channelFanoutJob{scope: channelFanoutMembers, channelID: channelID, pts: pts, recipients: []int64{42}, build: build} + } + d.Enqueue(context.Background(), job(7001, 1, func(context.Context, int64) *tg.Updates { + close(payloadStarted) + <-payloadRelease + return nil + })) + select { + case <-payloadStarted: + case <-time.After(time.Second): + cancel() + <-done + t.Fatal("payload worker did not enter blocking build") + } + d.Enqueue(context.Background(), job(7002, 2, func(context.Context, int64) *tg.Updates { return nil })) + // The next job cannot retain its payload and fills the sole shard overflow key. + d.Enqueue(context.Background(), job(fullChannel, 3, func(context.Context, int64) *tg.Updates { return nil })) + + // All key-bearing structures are now full. Repeated lost-channel producers must publish only + // a constant-size generation and return; no producer waiter or goroutine is allowed. + canceled, cancelRequest := context.WithCancel(context.Background()) + cancelRequest() + goroutinesBefore := runtime.NumGoroutine() + started := time.Now() + for pts := 4; pts <= 67; pts++ { + d.Enqueue(canceled, job(lostChannel, pts, func(context.Context, int64) *tg.Updates { return nil })) + } + elapsed := time.Since(started) + goroutinesAfter := runtime.NumGoroutine() + t.Logf("64 fully saturated Enqueue calls: elapsed=%v goroutines_before=%d goroutines_after=%d", elapsed, goroutinesBefore, goroutinesAfter) + if elapsed > 500*time.Millisecond { + close(payloadRelease) + cancel() + <-done + t.Fatalf("64 saturated Enqueue calls took %v; RPC producers must not wait for recovery capacity", elapsed) + } + if goroutinesAfter > goroutinesBefore+1 { + close(payloadRelease) + cancel() + <-done + t.Fatalf("producer goroutines grew from %d to %d; Enqueue must not spawn per-request waiters", goroutinesBefore, goroutinesAfter) + } + if generation := d.recoveryGeneration.Load(); generation == 0 { + close(payloadRelease) + cancel() + <-done + t.Fatal("all-memory saturation did not publish a recovery generation") + } + d.shards[0].mu.Lock() + waiters := d.shards[0].overflowWaiters + d.shards[0].mu.Unlock() + if waiters > 1 { + close(payloadRelease) + cancel() + <-done + t.Fatalf("overflow waiters = %d, want at most the one fixed recovery actor", waiters) + } + + // Restore capacity. The sweep no longer has the lost channel key in memory, so success proves + // it enumerated online membership and reloaded the authoritative max pts. + close(payloadRelease) + close(nudgeRelease) + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + pts := sessions.nudgePts(lostChannel) + if len(pts) > 0 && pts[len(pts)-1] == lostMaxPts && d.recoveryCompleted.Load() == d.recoveryGeneration.Load() { + break + } + time.Sleep(2 * time.Millisecond) + } + pts := sessions.nudgePts(lostChannel) + if len(pts) == 0 || pts[len(pts)-1] != lostMaxPts { + cancel() + <-done + t.Fatalf("lost channel nudges = %v, want durable max pts %d", pts, lostMaxPts) + } + if got, want := d.recoveryCompleted.Load(), d.recoveryGeneration.Load(); got != want { + cancel() + <-done + t.Fatalf("recovery generation completed=%d want=%d", got, want) + } + + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("dispatcher did not stop after saturated recovery") + } + if jobs, bytes := d.queuedBudgetSnapshot(); jobs != 0 || bytes != 0 { + t.Fatalf("shutdown queue budget = %d/%d, want 0/0", jobs, bytes) + } +} + +func TestChannelFanoutNudgeDeadlineRequestsRecoveryAndShutdownConverges(t *testing.T) { + blocked := make(chan struct{}) // deliberately never closed + sessions := newRecoveryFanoutSessions(nil, blocked) + channels := &recoveryFanoutChannels{pts: map[int64]int{}} + r := New(Config{}, Deps{Sessions: sessions, Channels: channels}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.nudgeTimeout = 25 * time.Millisecond + r.channelFanout = d + if !d.offerNudge(channelFanoutNudge{channelID: 8001, pts: 9}) { + t.Fatal("nudge rejected") + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + started := time.Now() + go func() { + d.Run(ctx) + close(done) + }() + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) && (sessions.pushErrors() == 0 || d.recoveryGeneration.Load() == 0) { + time.Sleep(time.Millisecond) + } + if sessions.pushErrors() == 0 || d.recoveryGeneration.Load() == 0 { + cancel() + <-done + t.Fatal("blocked nudge did not hit its explicit deadline and request recovery") + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + cancel() + <-done + t.Fatalf("nudge deadline observed after %v, want bounded worker occupancy", elapsed) + } + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("dispatcher shutdown did not cancel a blocked nudge push") + } +} + +func TestChannelFanoutShutdownCancelsCurrentlyBlockedNudgeWorker(t *testing.T) { + blocked := make(chan struct{}) // never released; only Run ctx may end the push + sessions := newRecoveryFanoutSessions(nil, blocked) + channels := &recoveryFanoutChannels{pts: map[int64]int{}} + r := New(Config{}, Deps{Sessions: sessions, Channels: channels}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.nudgeTimeout = time.Hour + r.channelFanout = d + if !d.offerNudge(channelFanoutNudge{channelID: 8101, pts: 10}) { + t.Fatal("nudge rejected") + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { d.Run(ctx); close(done) }() + select { + case <-sessions.pushStarted: + case <-time.After(time.Second): + cancel() + <-done + t.Fatal("nudge worker did not enter blocked push") + } + started := time.Now() + cancel() + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("Run did not converge after canceling a currently blocked nudge worker") + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("shutdown took %v with blocked nudge worker", elapsed) + } + if sessions.pushErrors() == 0 { + t.Fatal("blocked session push did not observe Run context cancellation") + } +} + +func TestChannelFanoutRecoveryRetainsGenerationOnErrorAndRepeatsConcurrentGeneration(t *testing.T) { + t.Run("max pts error is retried", func(t *testing.T) { + release := make(chan struct{}) + close(release) + sessions := newRecoveryFanoutSessions([]int64{10}, release) + channels := &recoveryFanoutChannels{pts: map[int64]int{10: 7}, failCalls: 1} + r := New(Config{}, Deps{Sessions: sessions, Channels: channels}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.nudgeWorkers = 0 + r.channelFanout = d + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { d.Run(ctx); close(done) }() + d.requestRecoverySweep("test injected max pts error") + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) && (channels.callCount() < 2 || d.recoveryCompleted.Load() != d.recoveryGeneration.Load()) { + time.Sleep(time.Millisecond) + } + if channels.callCount() < 2 { + t.Fatalf("MaxChannelPts calls = %d, want retry after error", channels.callCount()) + } + if got, want := d.recoveryCompleted.Load(), d.recoveryGeneration.Load(); got != want { + t.Fatalf("completed generation=%d want=%d after retry", got, want) + } + cancel() + <-done + }) + + t.Run("generation raised mid sweep forces a second full pass", func(t *testing.T) { + release := make(chan struct{}) + firstCalled := make(chan struct{}) + sessions := newRecoveryFanoutSessions([]int64{10}, nil) + channels := &recoveryFanoutChannels{pts: map[int64]int{10: 5}, firstCalled: firstCalled, release: release} + r := New(Config{}, Deps{Sessions: sessions, Channels: channels}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + d.nudgeWorkers = 0 + r.channelFanout = d + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { d.Run(ctx); close(done) }() + d.requestRecoverySweep("first generation") + select { + case <-firstCalled: + case <-time.After(time.Second): + cancel() + <-done + t.Fatal("first sweep did not reach MaxChannelPts") + } + channels.setPts(10, 9) + d.requestRecoverySweep("concurrent generation") + close(release) + deadline := time.Now().Add(time.Second) + for time.Now().Before(deadline) && (channels.callCount() < 2 || d.recoveryCompleted.Load() != d.recoveryGeneration.Load()) { + time.Sleep(time.Millisecond) + } + if channels.callCount() < 2 { + t.Fatalf("MaxChannelPts calls = %d, want two complete passes", channels.callCount()) + } + if got, want := d.recoveryCompleted.Load(), d.recoveryGeneration.Load(); got != want { + t.Fatalf("completed generation=%d want=%d", got, want) + } + select { + case channelID := <-d.nudgeJobs: + nudge, ok := d.takeNudge(channelID) + if !ok || nudge.channelID != 10 || nudge.pts != 9 { + t.Fatalf("coalesced nudge = %+v ok=%v, want channel=10 highest pts=9", nudge, ok) + } + case <-time.After(time.Second): + t.Fatal("coalesced recovery nudge was not queued") + } + cancel() + <-done + }) +} + +func TestChannelFanoutCoalescingNudgeMailboxKeepsHighestPts(t *testing.T) { + r := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System) + d := newChannelFanoutDispatcher(r, 1, 1) + for _, pts := range []int{5, 11, 7} { + if !d.offerNudge(channelFanoutNudge{channelID: 1001, pts: pts}) { + t.Fatalf("offer pts %d rejected", pts) + } + } + if got := len(d.nudgeJobs); got != 1 { + t.Fatalf("nudge queue cardinality = %d, want one slot for a hot channel", got) + } + channelID := <-d.nudgeJobs + nudge, ok := d.takeNudge(channelID) + if !ok || nudge.pts != 11 { + t.Fatalf("coalesced nudge = %+v ok=%v, want highest pts=11", nudge, ok) + } +} + +func TestChannelFanoutOverflowWaitHonorsContextStopAndHasNoLossyTimeout(t *testing.T) { + t.Run("context cancellation", func(t *testing.T) { + shard := newChannelFanoutShard(1) + shard.overflowLimit = 1 + if !shard.enqueueOverflow(1001, 5) { + t.Fatal("initial overflow watermark rejected") + } + ctx, cancel := context.WithCancel(context.Background()) + result := make(chan bool, 1) + go func() { result <- shard.enqueueOverflowWait(ctx, 2001, 7, make(chan struct{})) }() + cancel() + select { + case accepted := <-result: + if accepted { + t.Fatal("overflow wait accepted without mailbox space") + } + case <-time.After(time.Second): + t.Fatal("overflow wait ignored request context cancellation") + } + }) + + t.Run("dispatcher stop", func(t *testing.T) { + shard := newChannelFanoutShard(1) + shard.overflowLimit = 1 + if !shard.enqueueOverflow(1001, 5) { + t.Fatal("initial overflow watermark rejected") + } + stop := make(chan struct{}) + result := make(chan bool, 1) + go func() { result <- shard.enqueueOverflowWait(context.Background(), 2001, 7, stop) }() + close(stop) + select { + case accepted := <-result: + if accepted { + t.Fatal("overflow wait accepted without mailbox space") + } + case <-time.After(time.Second): + t.Fatal("overflow wait ignored dispatcher stop") + } + }) + + t.Run("waits beyond old maximum until space", func(t *testing.T) { + shard := newChannelFanoutShard(1) + shard.overflowLimit = 1 + if !shard.enqueueOverflow(1001, 5) { + t.Fatal("initial overflow watermark rejected") + } + result := make(chan bool, 1) + go func() { + result <- shard.enqueueOverflowWait(context.Background(), 2001, 7, make(chan struct{})) + }() + select { + case accepted := <-result: + t.Fatalf("overflow wait ended at the old fixed timeout: accepted=%v", accepted) + case <-time.After(150 * time.Millisecond): + } + if channelID, pts, ok := shard.popOverflow(); !ok || channelID != 1001 || pts != 5 { + t.Fatalf("released watermark = %d/%d/%v, want 1001/5/true", channelID, pts, ok) + } + select { + case accepted := <-result: + if !accepted { + t.Fatal("overflow wait rejected after mailbox space was released") + } + case <-time.After(time.Second): + t.Fatal("overflow wait did not resume after mailbox space was released") + } + }) +} diff --git a/internal/rpc/channels_legacy_chat.go b/internal/rpc/channels_legacy_chat.go index 346e135a..3174ec1c 100644 --- a/internal/rpc/channels_legacy_chat.go +++ b/internal/rpc/channels_legacy_chat.go @@ -445,14 +445,13 @@ func (r *Router) onMessagesSetChatTheme(ctx context.Context, req *tg.MessagesSet if err != nil { return nil, err } - authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) res, err := r.deps.Messages.SetChatTheme(ctx, userID, domain.SetPrivateChatThemeRequest{ OwnerUserID: userID, Peer: peer, Emoticon: emoticon, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) @@ -578,6 +577,7 @@ func (r *Router) enqueueChannelWallpaperFanout(ctx context.Context, originUserID fanoutCache := newViewerPeerCache(r) ownerIDs := channelMessageFanoutOwnerIDs(sendRes, nil) r.enqueueChannelFanoutWithPrefetch(ctx, channelFanoutMembers, originUserID, res.Channel.ID, res.Event.Pts, res.Recipients, + 0, func(bgCtx context.Context, viewers []int64) { r.prefetchChannelFanoutUsers(bgCtx, fanoutCache, viewers, ownerIDs) }, diff --git a/internal/rpc/channels_members.go b/internal/rpc/channels_members.go index 2c89a593..0deb431f 100644 --- a/internal/rpc/channels_members.go +++ b/internal/rpc/channels_members.go @@ -369,7 +369,11 @@ func (r *Router) recordChannelStateForUser(ctx context.Context, userID, channelI authKeyID, _ = AuthKeyIDFrom(ctx) excludeSessionID, _ = SessionIDFrom(ctx) } - event, _, err := r.deps.Updates.RecordChannelState(ctx, authKeyID, userID, channelID, excludeSessionID) + excludeAuthKeyID := [8]byte{} + if excludeCurrent { + excludeAuthKeyID = rawAuthKeyIDForOrigin(ctx) + } + event, _, err := r.deps.Updates.RecordChannelState(ctx, authKeyID, userID, channelID, excludeAuthKeyID, excludeSessionID) if err != nil { return } diff --git a/internal/rpc/channels_stubs.go b/internal/rpc/channels_stubs.go index 375d1a59..f6d7a36a 100644 --- a/internal/rpc/channels_stubs.go +++ b/internal/rpc/channels_stubs.go @@ -264,7 +264,7 @@ func (r *Router) recordChannelAvailableMessages(ctx context.Context, userID, cha } authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - recorded, _, err := r.deps.Updates.RecordChannelAvailableMessages(ctx, authKeyID, userID, channelID, availableMinID, sessionID) + recorded, _, err := r.deps.Updates.RecordChannelAvailableMessages(ctx, authKeyID, userID, channelID, availableMinID, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return event } @@ -300,7 +300,7 @@ func (r *Router) recordChannelReadInbox(ctx context.Context, userID int64, read StillUnreadCount: read.StillUnreadCount, ChannelPts: read.Pts, Changed: read.Changed, - }, sessionID) + }, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return domain.UpdateEvent{}, internalErr() } @@ -673,6 +673,8 @@ func channelInvalidErr(err error) error { return tgerr400("USER_ALREADY_PARTICIPANT") case errors.Is(err, domain.ErrReplyMessageIDInvalid): return replyMessageIDInvalidErr() + case errors.Is(err, domain.ErrMessageRandomIDDuplicate): + return randomIDDuplicateErr() default: if seconds, ok := domain.SlowModeWaitSeconds(err); ok { return tgerr.New(420, fmt.Sprintf("SLOWMODE_WAIT_%d", seconds)) diff --git a/internal/rpc/channels_topics.go b/internal/rpc/channels_topics.go index 723f0b20..cfbdb818 100644 --- a/internal/rpc/channels_topics.go +++ b/internal/rpc/channels_topics.go @@ -42,7 +42,7 @@ func (r *Router) onChannelsToggleViewForumAsMessages(ctx context.Context, req *t if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err = r.deps.Updates.RecordChannelViewForumAsMessages(ctx, authKeyID, userID, channelID, req.Enabled, sessionID) + event, _, err = r.deps.Updates.RecordChannelViewForumAsMessages(ctx, authKeyID, userID, channelID, req.Enabled, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return nil, internalErr() } diff --git a/internal/rpc/channels_updates.go b/internal/rpc/channels_updates.go index 74a48a44..48aad29f 100644 --- a/internal/rpc/channels_updates.go +++ b/internal/rpc/channels_updates.go @@ -206,6 +206,14 @@ func (r *Router) channelMessagesUpdatesWithPeerCache(ctx context.Context, viewer updates = append(updates, update) } } + if res.Duplicate && res.ReplayDeleteEvent != nil { + if update := tgChannelUpdate(viewerUserID, *res.ReplayDeleteEvent); update != nil { + updates = append(updates, update) + } + if res.ReplayDeleteEvent.Date > date { + date = res.ReplayDeleteEvent.Date + } + } collectChannelUpdatePeerRefs(res.Event, res.Channel.ID, userIDs, channelIDs) collectChannelMessagePeerRefs(res.Message, res.Channel.ID, userIDs, channelIDs) if date == 0 { diff --git a/internal/rpc/chatlists.go b/internal/rpc/chatlists.go index eb7e9647..c9a78320 100644 --- a/internal/rpc/chatlists.go +++ b/internal/rpc/chatlists.go @@ -379,7 +379,7 @@ func (r *Router) recordChatlistFilterUpdate(ctx context.Context, userID int64, f if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, sessionID) + event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return internalErr() } @@ -400,7 +400,7 @@ func (r *Router) chatlistFilterUpdates(ctx context.Context, userID int64, filter if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, sessionID) + event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, filterID, folder, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return nil, internalErr() } diff --git a/internal/rpc/contacts.go b/internal/rpc/contacts.go index 10794bd8..38bd6ebc 100644 --- a/internal/rpc/contacts.go +++ b/internal/rpc/contacts.go @@ -984,10 +984,10 @@ func (r *Router) recordAcceptedContactTargetUpdates(ctx context.Context, userID, return internalErr() } var zeroAuthKeyID [8]byte - if err := r.recordPeerSettingsForUser(ctx, zeroAuthKeyID, targetUserID, peer, settings, 0); err != nil { + if err := r.recordPeerSettingsForUser(ctx, zeroAuthKeyID, targetUserID, peer, settings, zeroAuthKeyID, 0); err != nil { return internalErr() } - if err := r.recordContactsResetForUser(ctx, zeroAuthKeyID, targetUserID, 0); err != nil { + if err := r.recordContactsResetForUser(ctx, zeroAuthKeyID, targetUserID, zeroAuthKeyID, 0); err != nil { return internalErr() } peerUser := domain.User{ID: userID} @@ -1017,14 +1017,14 @@ func (r *Router) pushContactsReset(ctx context.Context, userID int64) { func (r *Router) recordContactsReset(ctx context.Context, userID int64) error { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - return r.recordContactsResetForUser(ctx, authKeyID, userID, sessionID) + return r.recordContactsResetForUser(ctx, authKeyID, userID, rawAuthKeyIDForOrigin(ctx), sessionID) } -func (r *Router) recordContactsResetForUser(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) error { +func (r *Router) recordContactsResetForUser(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) error { if r.deps.Updates == nil || userID == 0 { return nil } - event, _, err := r.deps.Updates.RecordContactsReset(ctx, authKeyID, userID, excludeSessionID) + event, _, err := r.deps.Updates.RecordContactsReset(ctx, stateAuthKeyID, userID, excludeAuthKeyID, excludeSessionID) if err == nil && excludeSessionID != 0 { r.bookkeepAuxPtsForCurrentSession(ctx, event) } @@ -1034,7 +1034,7 @@ func (r *Router) recordContactsResetForUser(ctx context.Context, authKeyID [8]by func (r *Router) recordPeerSettings(ctx context.Context, userID int64, peer domain.Peer, settings domain.PeerSettings) error { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - return r.recordPeerSettingsForUser(ctx, authKeyID, userID, peer, settings, sessionID) + return r.recordPeerSettingsForUser(ctx, authKeyID, userID, peer, settings, rawAuthKeyIDForOrigin(ctx), sessionID) } func (r *Router) recordPeerStoryBlocked(ctx context.Context, userID int64, peer domain.Peer, blocked bool) error { @@ -1043,18 +1043,18 @@ func (r *Router) recordPeerStoryBlocked(ctx context.Context, userID int64, peer if r.deps.Updates == nil || userID == 0 { return nil } - event, _, err := r.deps.Updates.RecordPeerStoryBlocked(ctx, authKeyID, userID, peer, blocked, sessionID) + event, _, err := r.deps.Updates.RecordPeerStoryBlocked(ctx, authKeyID, userID, peer, blocked, rawAuthKeyIDForOrigin(ctx), sessionID) if err == nil && sessionID != 0 { r.bookkeepAuxPtsForCurrentSession(ctx, event) } return err } -func (r *Router) recordPeerSettingsForUser(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeSessionID int64) error { +func (r *Router) recordPeerSettingsForUser(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeAuthKeyID [8]byte, excludeSessionID int64) error { if r.deps.Updates == nil || userID == 0 { return nil } - event, _, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, settings, excludeSessionID) + event, _, err := r.deps.Updates.RecordPeerSettings(ctx, stateAuthKeyID, userID, peer, settings, excludeAuthKeyID, excludeSessionID) if err == nil && excludeSessionID != 0 { r.bookkeepAuxPtsForCurrentSession(ctx, event) } diff --git a/internal/rpc/context.go b/internal/rpc/context.go index 29846287..a86de795 100644 --- a/internal/rpc/context.go +++ b/internal/rpc/context.go @@ -16,8 +16,21 @@ const ( sessionIDKey userIDKey invokeWithoutUpdatesKey + inboundRPCBytesKey ) +func withInboundRPCBytes(ctx context.Context, n int) context.Context { + if n < 0 { + n = 0 + } + return context.WithValue(ctx, inboundRPCBytesKey, n) +} + +func inboundRPCBytesFrom(ctx context.Context) int { + v, _ := ctx.Value(inboundRPCBytesKey).(int) + return v +} + const currentClientLayer = 227 var androidSDKVersionRE = regexp.MustCompile(`\bsdk\s+\d+\b`) @@ -153,6 +166,16 @@ func AuthKeyIDFrom(ctx context.Context) ([8]byte, bool) { return v, ok } +// rawAuthKeyIDForOrigin 返回用于 update/outbox 当前 session 排除的物理 raw key。 +// 单测/非 edge 调用若没有注入 raw key，才回退业务 key；生产 Router context 两者都有。 +func rawAuthKeyIDForOrigin(ctx context.Context) [8]byte { + if id, ok := RawAuthKeyIDFrom(ctx); ok { + return id + } + id, _ := AuthKeyIDFrom(ctx) + return id +} + // WithSessionID 在 ctx 注入调用方的 MTProto session_id。 func WithSessionID(ctx context.Context, id int64) context.Context { return context.WithValue(ctx, sessionIDKey, id) diff --git a/internal/rpc/deps.go b/internal/rpc/deps.go index 74a28250..78769c05 100644 --- a/internal/rpc/deps.go +++ b/internal/rpc/deps.go @@ -71,6 +71,7 @@ type ScopedSessionBinder interface { UserIDResolvedForAuthKey(rawAuthKeyID [8]byte, sessionID int64) (userID int64, resolved bool) SetReceivesUpdatesForAuthKey(rawAuthKeyID [8]byte, sessionID int64, receives bool) PushToSessionForAuthKey(ctx context.Context, rawAuthKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error + // excludeAuthKeyID is the physical/raw auth key, paired with session_id. PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) } @@ -158,6 +159,21 @@ type ChannelNudgeProvider interface { OnlineChannelMemberUserIDsExcluding(channelID int64, exclude map[int64]struct{}, limit int) []int64 } +// ChannelFanoutRecoverySessionProvider snapshots the process-local online joined-channel index in +// stable channel-id order. It is used only after every keyed fan-out mailbox is saturated. The +// fixed recovery actor accepts the temporary 8*C id slice so one sweep never repeatedly scans the +// SessionManager index or holds its global lock while sorting/database work runs. +type ChannelFanoutRecoverySessionProvider interface { + OnlineChannelIDsSnapshot() []int64 +} + +// ChannelFanoutRecoveryPtsProvider reloads the authoritative channel pts after in-memory fan-out +// saturation. Production channels.Service implements it through the channel store; keeping this +// separate from ChannelsService avoids burdening lightweight RPC fakes that never run the worker. +type ChannelFanoutRecoveryPtsProvider interface { + MaxChannelPtsBatch(ctx context.Context, channelIDs []int64) (map[int64]int, error) +} + // RateLimiter 抽象 RPC 高频写操作限流。 type RateLimiter interface { Allow(ctx context.Context, key string, limit int, window time.Duration) (allowed bool, retryAfterSeconds int, err error) @@ -255,7 +271,7 @@ type UserPremiumStatusService interface { // AccountService 抽象账号设置查询。 type AccountService interface { SendChangePhoneCode(ctx context.Context, userID int64, authKeyID [8]byte, sessionID int64, phone string) (string, domain.AuthCodeDelivery, error) - ChangePhone(ctx context.Context, userID int64, authKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error) + ChangePhone(ctx context.Context, userID int64, authKeyID, originRawAuthKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error) GetPassword(ctx context.Context, userID int64) (domain.PasswordSettings, error) GetPasswordSettings(ctx context.Context, userID int64, check domain.PasswordCheck) (domain.PrivatePasswordSettings, error) UpdatePasswordSettings(ctx context.Context, userID int64, check domain.PasswordCheck, input domain.PasswordInputSettings) error @@ -270,10 +286,9 @@ type AccountService interface { SendLoginEmailCode(ctx context.Context, userID int64, phone, phoneCodeHash, email string, setup bool) (string, int, error) VerifyLoginEmail(ctx context.Context, userID int64, phone, phoneCodeHash, code string, setup bool) (string, error) SetLoginEmail(ctx context.Context, userID int64, email string) error - SetLoginEmailByPhone(ctx context.Context, phone, email string) error LoginEmail(ctx context.Context, userID int64) (string, bool, error) LoginEmailByPhone(ctx context.Context, phone string) (string, bool, error) - ClearLoginEmailByPhone(ctx context.Context, phone string) error + ClearLoginEmail(ctx context.Context, userID int64) error ResetPassword(ctx context.Context, userID int64) (domain.PasswordResetResult, error) DeclinePasswordReset(ctx context.Context, userID int64) error SaveMusic(ctx context.Context, userID int64, req domain.SaveMusicRequest) (bool, error) @@ -336,30 +351,30 @@ type UpdatesService interface { ClearAuthKey(ctx context.Context, authKeyID [8]byte) error RecordNewMessage(ctx context.Context, authKeyID [8]byte, userID int64, msg domain.Message) (domain.UpdateEvent, domain.UpdateState, error) PublishNewMessage(ctx context.Context, userID int64, msg domain.Message) (domain.UpdateEvent, domain.UpdateState, error) - RecordStory(ctx context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordStory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, story domain.Story, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) RecordStoryFanout(ctx context.Context, userID int64, story domain.Story) (domain.UpdateEvent, domain.UpdateState, error) - RecordReadStories(ctx context.Context, authKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordSentStoryReaction(ctx context.Context, authKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordQuickReplyMutation(ctx context.Context, authKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordReadHistory(ctx context.Context, authKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordContactsReset(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordChannelState(ctx context.Context, authKeyID [8]byte, userID, channelID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordPinnedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordSavedDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordPinnedSavedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordDialogUnreadMark(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordPeerSettings(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordPeerStoryBlocked(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordDialogFilter(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordDialogFilterOrder(ctx context.Context, authKeyID [8]byte, userID int64, order []int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordDialogFiltersReload(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordFolderPeers(ctx context.Context, authKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordChannelAvailableMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, availableMinID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordChannelViewForumAsMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, enabled bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordChannelDiscussionInbox(ctx context.Context, authKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) - RecordDraftMessage(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordReadStories(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordSentStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordNewStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordQuickReplyMutation(ctx context.Context, stateAuthKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordReadHistory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordContactsReset(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordChannelState(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + 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) + RecordPinnedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordSavedDialogPinned(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordPinnedSavedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordDialogUnreadMark(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + 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) + RecordPeerStoryBlocked(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordDialogFilter(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordDialogFilterOrder(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordDialogFiltersReload(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordFolderPeers(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordChannelAvailableMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, availableMinID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordChannelViewForumAsMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, enabled bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordChannelDiscussionInbox(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) + RecordDraftMessage(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) } // ContactsService 抽象通讯录查询。 @@ -455,6 +470,13 @@ type MessagesService interface { DeleteSavedHistory(ctx context.Context, userID int64, req domain.DeleteSavedHistoryRequest) (domain.DeleteSavedHistoryResult, error) } +// AlbumGroupService 是 MessagesService 的可选、生产必备能力：sendMultiMedia 在 +// 解析任何媒体或落第一条消息前，持久预留整批 random_id 的 grouped_id。 +// 单独定义可避免让不触发 sendMultiMedia 的轻量测试替身实现无关方法。 +type AlbumGroupService interface { + ReserveAlbumGroup(ctx context.Context, userID int64, req domain.AlbumGroupReservationRequest) (int64, error) +} + // StoriesService 抽象 story 读取、已读、观看与 reaction 状态。 type StoriesService interface { CreateStory(ctx context.Context, userID int64, req domain.StoryCreateRequest) (domain.StoryCreateResult, error) diff --git a/internal/rpc/dialogs_rpc_test.go b/internal/rpc/dialogs_rpc_test.go index e7f59425..224dc516 100644 --- a/internal/rpc/dialogs_rpc_test.go +++ b/internal/rpc/dialogs_rpc_test.go @@ -172,7 +172,11 @@ func TestMessagesGetPeerDialogsReturnsRequestedDialogsAndState(t *testing.T) { func TestDialogSettingRPCsRecordDurableUpdates(t *testing.T) { var authKeyID [8]byte authKeyID[0] = 9 - ctx := WithSessionID(WithAuthKeyID(WithUserID(context.Background(), 1000000001), authKeyID), 55) + rawAuthKeyID := [8]byte{9, 7} + ctx := WithRawAuthKeyID( + WithSessionID(WithAuthKeyID(WithUserID(context.Background(), 1000000001), authKeyID), 55), + rawAuthKeyID, + ) peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002} dialogPeer := &tg.InputDialogPeer{Peer: &tg.InputPeerUser{UserID: peer.ID}} dialogs := &captureDialogs{} @@ -187,6 +191,9 @@ func TestDialogSettingRPCsRecordDurableUpdates(t *testing.T) { if len(updates.events) != 1 || updates.events[0].Type != domain.UpdateEventDialogPinned || updates.events[0].Peer != peer || !updates.events[0].Bool || updates.excludeSessionID != 55 { t.Fatalf("pin event = %+v, want durable dialog_pinned", updates.events) } + if updates.authKeyID != authKeyID || updates.excludeAuthKeyID != rawAuthKeyID { + t.Fatalf("durable update keys = state:%x exclude:%x, want business:%x raw:%x", updates.authKeyID, updates.excludeAuthKeyID, authKeyID, rawAuthKeyID) + } if ok, err := r.onMessagesReorderPinnedDialogs(ctx, &tg.MessagesReorderPinnedDialogsRequest{Order: []tg.InputDialogPeerClass{dialogPeer}}); err != nil || !ok { t.Fatalf("reorder pinned = %v, %v", ok, err) diff --git a/internal/rpc/encrypted_chats.go b/internal/rpc/encrypted_chats.go index 9c2779c9..fca5532f 100644 --- a/internal/rpc/encrypted_chats.go +++ b/internal/rpc/encrypted_chats.go @@ -88,7 +88,7 @@ func (r *Router) recordEncryptionEventBestEffort(ctx context.Context, chatID int // 全部活跃密聊，并向对端推送 encryptedChatDiscarded（在线）+ 写 durable 事件（离线 getDifference // 补偿）。ownerUserID 是被销毁设备的所有者，用于定位对端。best-effort：失败仅记日志，绝不阻断 // 登出/撤销。修复 P1：此前 onAuthLogOut 等不级联 discard，对端继续往死 auth_key 投递成静默死链 -//（消息 acked=f / qts 永久积压，对端永看不到 discarded）。 +// （消息 acked=f / qts 永久积压，对端永看不到 discarded）。 func (r *Router) discardSecretChatsForAuthKey(ctx context.Context, businessAuthKeyID, ownerUserID int64) { if r.deps.SecretChats == nil || businessAuthKeyID == 0 || ownerUserID == 0 { return diff --git a/internal/rpc/errors.go b/internal/rpc/errors.go index a99d5c2e..57d05d07 100644 --- a/internal/rpc/errors.go +++ b/internal/rpc/errors.go @@ -271,6 +271,8 @@ func chatForwardsRestrictedErr() error { return tgerr.New(400, "CHAT_FORWARDS_RE func inputRequestInvalidErr() error { return tgerr.New(400, "INPUT_REQUEST_INVALID") } +func inputRequestTooLongErr() error { return tgerr.New(400, "INPUT_REQUEST_TOO_LONG") } + func persistentTimestampInvalidErr() error { return tgerr.New(400, "PERSISTENT_TIMESTAMP_INVALID") } func channelForumMissingErr() error { return tgerr.New(400, "CHANNEL_FORUM_MISSING") } @@ -281,6 +283,10 @@ func topicIDInvalidErr() error { return tgerr.New(400, "TOPIC_ID_INVALID") } // randomIDEmptyErr 表示发送消息缺少 random_id。 func randomIDEmptyErr() error { return tgerr.New(400, "RANDOM_ID_EMPTY") } +// randomIDDuplicateErr 表示同一发送者重复使用 random_id，但请求载荷与首次 +// 成功发送不一致。Layer 227 为该错误定义的 code 是 500。 +func randomIDDuplicateErr() error { return tgerr.New(500, "RANDOM_ID_DUPLICATE") } + // scheduleDateInvalidErr 表示当前阶段不支持定时消息。 func scheduleDateInvalidErr() error { return tgerr.New(400, "SCHEDULE_DATE_INVALID") } diff --git a/internal/rpc/folders.go b/internal/rpc/folders.go index 7fc60d61..5e6babf4 100644 --- a/internal/rpc/folders.go +++ b/internal/rpc/folders.go @@ -53,7 +53,7 @@ func (r *Router) onFoldersEditPeerFolders(ctx context.Context, folderPeers []tg. if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err = r.deps.Updates.RecordFolderPeers(ctx, authKeyID, userID, updates, sessionID) + event, _, err = r.deps.Updates.RecordFolderPeers(ctx, authKeyID, userID, updates, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return nil, internalErr() } diff --git a/internal/rpc/idle_backoff.go b/internal/rpc/idle_backoff.go index e6e123b5..7faa5f4a 100644 --- a/internal/rpc/idle_backoff.go +++ b/internal/rpc/idle_backoff.go @@ -57,7 +57,10 @@ func runIdleBackoffLoop(ctx context.Context, interval, maxIdleInterval time.Dura case <-timer.C: } if dispatch(ctx) { - timer.Reset(backoff.ActiveDelay()) + // 有积压时立即继续 drain；interval 只用于空闲轮询。旧逻辑每个非空 + // batch 也固定等待 base，形成 batch/base 的人工吞吐上限。 + _ = backoff.ActiveDelay() + timer.Reset(0) continue } timer.Reset(backoff.IdleDelay()) diff --git a/internal/rpc/idle_backoff_test.go b/internal/rpc/idle_backoff_test.go index 39d83c30..02e9096b 100644 --- a/internal/rpc/idle_backoff_test.go +++ b/internal/rpc/idle_backoff_test.go @@ -1,6 +1,8 @@ package rpc import ( + "context" + "sync/atomic" "testing" "time" ) @@ -27,6 +29,35 @@ func TestIdleBackoffSequenceAndReset(t *testing.T) { } } +func TestIdleBackoffLoopDrainsActiveWorkImmediately(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + var calls atomic.Int32 + done := make(chan struct{}) + started := time.Now() + go func() { + defer close(done) + runIdleBackoffLoop(ctx, time.Second, time.Second, func(context.Context) bool { + if calls.Add(1) < 4 { + return true + } + cancel() + return false + }) + }() + select { + case <-done: + case <-time.After(300 * time.Millisecond): + t.Fatal("active drain waited for idle interval") + } + if got := calls.Load(); got != 4 { + t.Fatalf("dispatch calls = %d, want 4 consecutive active drains", got) + } + if elapsed := time.Since(started); elapsed >= 300*time.Millisecond { + t.Fatalf("active drain elapsed = %v, want no 1s base delay", elapsed) + } +} + func TestIdleBackoffSanitizesMaxBelowBase(t *testing.T) { backoff := newIdleBackoff(2*time.Second, time.Second) if got := backoff.IdleDelay(); got != 2*time.Second { diff --git a/internal/rpc/message_idempotency.go b/internal/rpc/message_idempotency.go new file mode 100644 index 00000000..c70677f8 --- /dev/null +++ b/internal/rpc/message_idempotency.go @@ -0,0 +1,84 @@ +package rpc + +import ( + "crypto/sha256" + "fmt" + + "github.com/gotd/td/bin" + "github.com/gotd/td/tg" +) + +// rpcRequestFingerprint 在任何自动实体补全、链接预览解析或上传媒体落库之前， +// 对客户端原始 TL request 取稳定指纹。这样 lost-response 重放不会因服务端派生 +// photo/document id、pending webpage 状态等变化而被误判为另一条消息。 +func rpcRequestFingerprint(req bin.Encoder) ([]byte, error) { + if req == nil { + return nil, fmt.Errorf("fingerprint rpc request: nil request") + } + var b bin.Buffer + if err := req.Encode(&b); err != nil { + return nil, fmt.Errorf("fingerprint rpc request: %w", err) + } + sum := sha256.Sum256(b.Raw()) + return sum[:], nil +} + +// sendMessageIdempotencyFingerprint fingerprints only the durable message intent. +// clear_draft/background/update_stickersets_order are one-shot client-side delivery +// hints: after a lost response DrKLO/TDesktop may legitimately retry the same +// random_id without them. They must not turn an exact send replay into +// RANDOM_ID_DUPLICATE. +func sendMessageIdempotencyFingerprint(req *tg.MessagesSendMessageRequest) ([]byte, error) { + if req == nil { + return nil, fmt.Errorf("fingerprint messages.sendMessage: nil request") + } + clone := *req + clone.Flags = 0 + clone.ClearDraft = false + clone.Background = false + clone.UpdateStickersetsOrder = false + return rpcRequestFingerprint(&clone) +} + +func sendMediaIdempotencyFingerprint(req *tg.MessagesSendMediaRequest) ([]byte, error) { + if req == nil { + return nil, fmt.Errorf("fingerprint messages.sendMedia: nil request") + } + clone := *req + clone.Flags = 0 + clone.ClearDraft = false + clone.Background = false + clone.UpdateStickersetsOrder = false + return rpcRequestFingerprint(&clone) +} + +// sendMultiMediaItemIdempotencyFingerprint deliberately reduces a batch to one +// InputSingleMedia. A retry containing only the failed subset therefore produces +// the same fingerprint for every surviving random_id as the original batch. +func sendMultiMediaItemIdempotencyFingerprint(req *tg.MessagesSendMultiMediaRequest, item tg.InputSingleMedia) ([]byte, error) { + if req == nil { + return nil, fmt.Errorf("fingerprint messages.sendMultiMedia item: nil request") + } + clone := *req + clone.Flags = 0 + clone.ClearDraft = false + clone.Background = false + clone.UpdateStickersetsOrder = false + clone.MultiMedia = []tg.InputSingleMedia{item} + return rpcRequestFingerprint(&clone) +} + +// forwardMessagesItemIdempotencyFingerprint makes the source message id and its +// paired random_id the unit of idempotency. Hashing the full ID/RandomID vectors +// incorrectly rejects a legal retry that contains only a failed subset. +func forwardMessagesItemIdempotencyFingerprint(req *tg.MessagesForwardMessagesRequest, messageID int, randomID int64) ([]byte, error) { + if req == nil { + return nil, fmt.Errorf("fingerprint messages.forwardMessages item: nil request") + } + clone := *req + clone.Flags = 0 + clone.Background = false + clone.ID = []int{messageID} + clone.RandomID = []int64{randomID} + return rpcRequestFingerprint(&clone) +} diff --git a/internal/rpc/message_idempotency_test.go b/internal/rpc/message_idempotency_test.go new file mode 100644 index 00000000..7633995b --- /dev/null +++ b/internal/rpc/message_idempotency_test.go @@ -0,0 +1,315 @@ +package rpc + +import ( + "bytes" + "context" + "fmt" + "testing" + + "github.com/gotd/td/clock" + "github.com/gotd/td/tg" + "github.com/gotd/td/tgerr" + "go.uber.org/zap/zaptest" + + "telesrv/internal/domain" +) + +func TestRPCRequestFingerprintStableAndPayloadSensitive(t *testing.T) { + req := &tg.MessagesSendMessageRequest{ + Peer: &tg.InputPeerUser{UserID: 1002, AccessHash: 22}, + Message: "hello", + RandomID: 991, + Entities: []tg.MessageEntityClass{&tg.MessageEntityBold{Offset: 0, Length: 5}}, + } + + first, err := rpcRequestFingerprint(req) + if err != nil { + t.Fatalf("first fingerprint: %v", err) + } + second, err := rpcRequestFingerprint(req) + if err != nil { + t.Fatalf("second fingerprint: %v", err) + } + if len(first) != 32 || !bytes.Equal(first, second) { + t.Fatalf("fingerprints = %x / %x, want stable SHA-256", first, second) + } + + req.Message = "changed" + changed, err := rpcRequestFingerprint(req) + if err != nil { + t.Fatalf("changed fingerprint: %v", err) + } + if bytes.Equal(first, changed) { + t.Fatalf("changed payload fingerprint = %x, want different from %x", changed, first) + } +} + +func TestSendMessageFingerprintIgnoresRetryOnlyHints(t *testing.T) { + first := &tg.MessagesSendMessageRequest{ + Peer: &tg.InputPeerUser{UserID: 1002, AccessHash: 22}, + Message: "hello", + RandomID: 991, + ClearDraft: true, + Background: true, + UpdateStickersetsOrder: true, + } + retry := *first + retry.Flags.Set(31) // stale decoded flags must not leak into the canonical intent. + retry.ClearDraft = false + retry.Background = false + retry.UpdateStickersetsOrder = false + + a, err := sendMessageIdempotencyFingerprint(first) + if err != nil { + t.Fatalf("first fingerprint: %v", err) + } + b, err := sendMessageIdempotencyFingerprint(&retry) + if err != nil { + t.Fatalf("retry fingerprint: %v", err) + } + if !bytes.Equal(a, b) { + t.Fatalf("retry-only flags changed fingerprint: %x != %x", a, b) + } + + retry.Message = "different" + c, err := sendMessageIdempotencyFingerprint(&retry) + if err != nil { + t.Fatalf("changed fingerprint: %v", err) + } + if bytes.Equal(a, c) { + t.Fatal("durable message change did not change fingerprint") + } +} + +func TestSendMultiMediaFingerprintIsPerItemAndSubsetStable(t *testing.T) { + item1 := tg.InputSingleMedia{Media: &tg.InputMediaEmpty{}, RandomID: 101, Message: "one"} + item2 := tg.InputSingleMedia{Media: &tg.InputMediaEmpty{}, RandomID: 102, Message: "two"} + full := &tg.MessagesSendMultiMediaRequest{ + Peer: &tg.InputPeerUser{UserID: 1002, AccessHash: 22}, + ClearDraft: true, + Background: true, + MultiMedia: []tg.InputSingleMedia{item1, item2}, + } + subset := &tg.MessagesSendMultiMediaRequest{ + Peer: full.Peer, + MultiMedia: []tg.InputSingleMedia{item2}, + } + + fromFull, err := sendMultiMediaItemIdempotencyFingerprint(full, item2) + if err != nil { + t.Fatalf("full item fingerprint: %v", err) + } + fromSubset, err := sendMultiMediaItemIdempotencyFingerprint(subset, item2) + if err != nil { + t.Fatalf("subset item fingerprint: %v", err) + } + if !bytes.Equal(fromFull, fromSubset) { + t.Fatalf("subset retry fingerprint = %x, want %x", fromSubset, fromFull) + } + + changed := item2 + changed.Message = "changed" + different, err := sendMultiMediaItemIdempotencyFingerprint(subset, changed) + if err != nil { + t.Fatalf("changed item fingerprint: %v", err) + } + if bytes.Equal(fromFull, different) { + t.Fatal("changed album item reused the original fingerprint") + } +} + +func TestForwardFingerprintIsPerItemAndSubsetStable(t *testing.T) { + full := &tg.MessagesForwardMessagesRequest{ + FromPeer: &tg.InputPeerUser{UserID: 1002, AccessHash: 22}, + ID: []int{41, 42}, + RandomID: []int64{201, 202}, + ToPeer: &tg.InputPeerUser{UserID: 1003, AccessHash: 33}, + Background: true, + } + subset := *full + subset.Flags = 0 + subset.Background = false + subset.ID = []int{42} + subset.RandomID = []int64{202} + + fromFull, err := forwardMessagesItemIdempotencyFingerprint(full, 42, 202) + if err != nil { + t.Fatalf("full item fingerprint: %v", err) + } + fromSubset, err := forwardMessagesItemIdempotencyFingerprint(&subset, 42, 202) + if err != nil { + t.Fatalf("subset item fingerprint: %v", err) + } + if !bytes.Equal(fromFull, fromSubset) { + t.Fatalf("subset retry fingerprint = %x, want %x", fromSubset, fromFull) + } + + different, err := forwardMessagesItemIdempotencyFingerprint(&subset, 41, 202) + if err != nil { + t.Fatalf("changed source fingerprint: %v", err) + } + if bytes.Equal(fromFull, different) { + t.Fatal("different source message reused the original fingerprint") + } +} + +func TestMessageSendErrMapsRandomIDConflict(t *testing.T) { + err := messageSendErr(fmt.Errorf("wrapped store error: %w", domain.ErrMessageRandomIDDuplicate)) + if !tgerr.Is(err, "RANDOM_ID_DUPLICATE") || !tgerr.IsCode(err, 500) { + t.Fatalf("messageSendErr = %v, want 500 RANDOM_ID_DUPLICATE", err) + } +} + +func TestMessageForwardErrMapsRandomIDConflict(t *testing.T) { + err := messageForwardErr(fmt.Errorf("wrapped store error: %w", domain.ErrMessageRandomIDDuplicate)) + if !tgerr.Is(err, "RANDOM_ID_DUPLICATE") || !tgerr.IsCode(err, 500) { + t.Fatalf("messageForwardErr = %v, want 500 RANDOM_ID_DUPLICATE", err) + } +} + +func TestPrivateSendDuplicateResponseIncludesAndroidConfirmationSnapshot(t *testing.T) { + res := domain.SendPrivateTextResult{ + Duplicate: true, + SenderMessage: domain.Message{ + ID: 41, UID: 51, RandomID: 9911, OwnerUserID: 1001, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: 1002}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: 1001}, + Date: 1700000000, Out: true, Body: "confirmed", Pts: 7, + }, + SenderEvent: domain.UpdateEvent{Pts: 7, PtsCount: 1, Date: 1700000000}, + } + updates := tgPrivateSendResultUpdates(res, 9911, true, nil, nil) + if len(updates.Updates) != 2 { + t.Fatalf("duplicate updates = %#v, want mapping + new message", updates.Updates) + } + mapping, ok := updates.Updates[0].(*tg.UpdateMessageID) + if !ok || mapping.ID != 41 || mapping.RandomID != 9911 { + t.Fatalf("duplicate mapping = %#v, want id/random_id 41/9911", updates.Updates[0]) + } + confirmed, ok := updates.Updates[1].(*tg.UpdateNewMessage) + if !ok || confirmed.Pts != 7 || confirmed.PtsCount != 1 { + t.Fatalf("duplicate confirmation = %#v, want UpdateNewMessage pts 7/1", updates.Updates[1]) + } + msg, ok := confirmed.Message.(*tg.Message) + if !ok || msg.ID != 41 || msg.Message != "confirmed" { + t.Fatalf("duplicate confirmation message = %#v, want sender snapshot", confirmed.Message) + } +} + +func TestPrivateSendDeletedDuplicateConfirmsThenConvergesWithDurableDelete(t *testing.T) { + deleteEvent := domain.UpdateEvent{UserID: 1001, Type: domain.UpdateEventDeleteMessages, Pts: 9, PtsCount: 1, Date: 1700000002, MessageIDs: []int{41}} + res := domain.SendPrivateTextResult{ + Duplicate: true, + SenderMessage: domain.Message{ + ID: 41, UID: 51, RandomID: 9911, OwnerUserID: 1001, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: 1002}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: 1001}, + Date: 1700000000, Out: true, Body: "original", Pts: 7, + }, + SenderEvent: domain.UpdateEvent{Pts: 7, PtsCount: 1, Date: 1700000000}, + ReplayDeleteEvent: &deleteEvent, + } + updates := tgPrivateSendResultUpdates(res, 9911, true, nil, nil) + if len(updates.Updates) != 3 { + t.Fatalf("deleted duplicate updates = %#v, want mapping + new + delete", updates.Updates) + } + if _, ok := updates.Updates[1].(*tg.UpdateNewMessage); !ok { + t.Fatalf("deleted duplicate confirmation = %#v, want UpdateNewMessage", updates.Updates[1]) + } + deleted, ok := updates.Updates[2].(*tg.UpdateDeleteMessages) + if !ok || deleted.Pts != 9 || deleted.PtsCount != 1 || len(deleted.Messages) != 1 || deleted.Messages[0] != 41 { + t.Fatalf("deleted duplicate convergence = %#v, want real delete event", updates.Updates[2]) + } +} + +func TestForwardDeletedDuplicateIncludesMessageAndDelete(t *testing.T) { + deleteEvent := &domain.UpdateEvent{UserID: 1001, Type: domain.UpdateEventDeleteMessages, Pts: 12, PtsCount: 1, Date: 1700000012, MessageIDs: []int{61}} + updates := tgForwardMessagesUpdates(domain.ForwardPrivateMessagesResult{ + SenderMessages: []domain.Message{{ + ID: 61, UID: 71, RandomID: 8811, OwnerUserID: 1001, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: 1002}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: 1001}, + Date: 1700000010, Out: true, Body: "forwarded", Pts: 10, + }}, + SenderEvents: []domain.UpdateEvent{{Pts: 10, PtsCount: 1, Date: 1700000010}}, + Duplicates: []bool{true}, + ReplayDeleteEvents: []*domain.UpdateEvent{deleteEvent}, + }, []int64{8811}, nil, nil) + if len(updates.Updates) != 3 { + t.Fatalf("forward duplicate updates = %#v, want mapping + new + delete", updates.Updates) + } + if _, ok := updates.Updates[1].(*tg.UpdateNewMessage); !ok { + t.Fatalf("forward duplicate confirmation = %#v, want UpdateNewMessage", updates.Updates[1]) + } + if deleted, ok := updates.Updates[2].(*tg.UpdateDeleteMessages); !ok || len(deleted.Messages) != 1 || deleted.Messages[0] != 61 || deleted.Pts != 12 { + t.Fatalf("forward duplicate delete = %#v, want durable delete", updates.Updates[2]) + } +} + +func TestChannelDeletedDuplicateEchoIncludesMessageAndDelete(t *testing.T) { + router := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System) + deleteEvent := &domain.ChannelUpdateEvent{ChannelID: 2001, Type: domain.ChannelUpdateDeleteMessages, Pts: 12, PtsCount: 1, Date: 1700000012, MessageIDs: []int{61}} + msg := domain.ChannelMessage{ + ChannelID: 2001, ID: 61, RandomID: 8811, SenderUserID: 1001, + From: domain.Peer{Type: domain.PeerTypeUser, ID: 1001}, + Date: 1700000010, Body: "channel", Pts: 10, + } + updates := router.channelMessagesUpdatesWithPeerCache(context.Background(), 1001, []domain.SendChannelMessageResult{{ + Channel: domain.Channel{ID: 2001, AccessHash: 22, Title: "group", Megagroup: true, Date: 1700000000}, + Message: msg, + Event: domain.ChannelUpdateEvent{ + ChannelID: 2001, Type: domain.ChannelUpdateNewMessage, Pts: 10, PtsCount: 1, Date: 1700000010, Message: msg, + }, + Duplicate: true, + ReplayDeleteEvent: deleteEvent, + }}, []int64{8811}, true, nil, newViewerPeerCache(router)) + if len(updates.Updates) != 3 { + t.Fatalf("channel duplicate updates = %#v, want mapping + new + delete", updates.Updates) + } + if _, ok := updates.Updates[1].(*tg.UpdateNewChannelMessage); !ok { + t.Fatalf("channel duplicate confirmation = %#v, want UpdateNewChannelMessage", updates.Updates[1]) + } + if deleted, ok := updates.Updates[2].(*tg.UpdateDeleteChannelMessages); !ok || deleted.ChannelID != 2001 || len(deleted.Messages) != 1 || deleted.Messages[0] != 61 || deleted.Pts != 12 { + t.Fatalf("channel duplicate delete = %#v, want durable channel delete", updates.Updates[2]) + } +} + +func TestPrivateSendRecordsRawOriginAuthKey(t *testing.T) { + const ( + senderID = int64(1001) + recipientID = int64(1002) + ) + raw := [8]byte{1, 2, 3, 4, 5, 6, 7, 8} + business := [8]byte{8, 7, 6, 5, 4, 3, 2, 1} + messages := &captureMessages{} + router := New(Config{}, Deps{ + Messages: messages, + Users: mapUsersService{users: map[int64]domain.User{ + senderID: {ID: senderID, FirstName: "Sender"}, + recipientID: {ID: recipientID, FirstName: "Recipient"}, + }}, + }, zaptest.NewLogger(t), clock.System) + ctx := WithSessionID( + WithRawAuthKeyID( + WithAuthKeyID( + WithUserID(context.Background(), senderID), + business, + ), + raw, + ), + 77, + ) + if _, err := router.onMessagesSendMessage(ctx, &tg.MessagesSendMessageRequest{ + Peer: &tg.InputPeerUser{UserID: recipientID}, + Message: "raw origin", + RandomID: 992, + }); err != nil { + t.Fatalf("send message: %v", err) + } + if messages.sendReq.OriginAuthKeyID != raw || messages.sendReq.OriginAuthKeyID == business { + t.Fatalf("origin auth key = %x, want raw %x (business %x)", messages.sendReq.OriginAuthKeyID, raw, business) + } + if messages.sendReq.OriginSessionID != 77 { + t.Fatalf("origin session = %d, want 77", messages.sendReq.OriginSessionID) + } +} diff --git a/internal/rpc/message_replay_preflight_test.go b/internal/rpc/message_replay_preflight_test.go new file mode 100644 index 00000000..6d1d33b2 --- /dev/null +++ b/internal/rpc/message_replay_preflight_test.go @@ -0,0 +1,454 @@ +package rpc + +import ( + "context" + "reflect" + "testing" + "time" + + "github.com/gotd/td/clock" + "github.com/gotd/td/tg" + "go.uber.org/zap/zaptest" + + appchannels "telesrv/internal/app/channels" + appusers "telesrv/internal/app/users" + "telesrv/internal/domain" + "telesrv/internal/store/memory" +) + +type replayPreflightMessages struct { + *captureMessages + replays map[int64]domain.SendPrivateTextResult + lookupRequests []domain.PrivateSendReplayRequest + sendRequests []domain.SendPrivateTextRequest + reserveRequests []domain.AlbumGroupReservationRequest + reservedGroupedID int64 + forceSendDuplicate bool +} + +func newReplayPreflightMessages() *replayPreflightMessages { + return &replayPreflightMessages{ + captureMessages: &captureMessages{}, + replays: make(map[int64]domain.SendPrivateTextResult), + reservedGroupedID: 81001, + } +} + +func (s *replayPreflightMessages) LookupPrivateSendReplay(_ context.Context, _ int64, req domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) { + s.lookupRequests = append(s.lookupRequests, req) + res, found := s.replays[req.RandomID] + if found { + res.Duplicate = true + } + return res, found, nil +} + +func (s *replayPreflightMessages) SendPrivateText(_ context.Context, _ int64, req domain.SendPrivateTextRequest) (domain.SendPrivateTextResult, error) { + s.sendRequests = append(s.sendRequests, req) + res := privateReplayFixture(req.SenderUserID, req.RecipientUserID, req.RandomID, 100+len(s.sendRequests), req.GroupedID) + res.Duplicate = s.forceSendDuplicate + s.replays[req.RandomID] = res + return res, nil +} + +func (s *replayPreflightMessages) ReserveAlbumGroup(_ context.Context, _ int64, req domain.AlbumGroupReservationRequest) (int64, error) { + s.reserveRequests = append(s.reserveRequests, req) + return s.reservedGroupedID, nil +} + +func privateReplayFixture(senderID, recipientID, randomID int64, messageID int, groupedID int64) domain.SendPrivateTextResult { + msg := domain.Message{ + ID: messageID, + UID: int64(messageID) + 1000, + RandomID: randomID, + OwnerUserID: senderID, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: recipientID}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: senderID}, + Date: 1700000000 + messageID, + Out: true, + Body: "committed", + GroupedID: groupedID, + Pts: messageID, + } + event := domain.UpdateEvent{ + UserID: senderID, + Type: domain.UpdateEventNewMessage, + Pts: messageID, + PtsCount: 1, + Date: msg.Date, + Message: msg, + } + return domain.SendPrivateTextResult{ + SenderMessage: msg, + SenderEvent: event, + } +} + +type replayCountingDialogs struct { + *captureDialogs + deleteDraftCalls int +} + +func (s *replayCountingDialogs) DeleteDraft(_ context.Context, _ int64, peer domain.Peer, topMessageID int) (bool, error) { + s.deleteDraftCalls++ + s.deletedDraft.peer = peer + s.deletedDraft.topMessageID = topMessageID + return true, nil +} + +type replaySelectiveFiles struct { + *fakeFiles + getPhotoIDs []int64 +} + +func (f *replaySelectiveFiles) GetPhoto(ctx context.Context, id int64) (domain.Photo, bool, error) { + f.getPhotoIDs = append(f.getPhotoIDs, id) + return f.fakeFiles.GetPhoto(ctx, id) +} + +func TestSendMessageExactReplayPrecedesSaturatedRateLimiter(t *testing.T) { + const userID = int64(1001) + messages := newReplayPreflightMessages() + messages.replays[7001] = privateReplayFixture(userID, userID, 7001, 71, 0) + limiter := &captureRateLimiter{block: true, retryAfter: 30} + dialogs := &replayCountingDialogs{captureDialogs: &captureDialogs{}} + r := New(Config{SendRateLimit: 1, SendRateWindow: time.Minute}, Deps{ + Messages: messages, + Dialogs: dialogs, + Limiter: limiter, + }, zaptest.NewLogger(t), clock.System) + + updates, err := r.onMessagesSendMessage(WithUserID(context.Background(), userID), &tg.MessagesSendMessageRequest{ + Peer: &tg.InputPeerSelf{}, + Message: "already committed", + RandomID: 7001, + ClearDraft: true, + }) + if err != nil { + t.Fatalf("exact replay: %v", err) + } + if _, ok := updates.(*tg.Updates); !ok { + t.Fatalf("updates = %T, want *tg.Updates", updates) + } + if len(limiter.calls) != 0 { + t.Fatalf("limiter calls = %+v, want none for exact replay", limiter.calls) + } + if len(messages.sendRequests) != 0 { + t.Fatalf("send requests = %d, want 0", len(messages.sendRequests)) + } + if dialogs.deleteDraftCalls != 0 { + t.Fatalf("draft deletes = %d, want 0 for duplicate", dialogs.deleteDraftCalls) + } +} + +func TestSendMessageConcurrentReplayRaceDoesNotClearDraft(t *testing.T) { + const userID = int64(1001) + messages := newReplayPreflightMessages() + // The read-only preflight misses, then the atomic store fence observes that another request + // committed the same random_id first and returns Duplicate=true. + messages.forceSendDuplicate = true + dialogs := &replayCountingDialogs{captureDialogs: &captureDialogs{}} + r := New(Config{}, Deps{Messages: messages, Dialogs: dialogs}, zaptest.NewLogger(t), clock.System) + + if _, err := r.onMessagesSendMessage(WithUserID(context.Background(), userID), &tg.MessagesSendMessageRequest{ + Peer: &tg.InputPeerSelf{}, + Message: "concurrent exact replay", + RandomID: 7002, + ClearDraft: true, + }); err != nil { + t.Fatalf("concurrent replay race: %v", err) + } + if len(messages.lookupRequests) != 1 || len(messages.sendRequests) != 1 { + t.Fatalf("lookup/send calls = %d/%d, want preflight miss then one atomic send", len(messages.lookupRequests), len(messages.sendRequests)) + } + if dialogs.deleteDraftCalls != 0 { + t.Fatalf("draft deletes = %d, want 0 when store race returns duplicate", dialogs.deleteDraftCalls) + } +} + +func TestSendMediaExactReplayPrecedesMediaResolvers(t *testing.T) { + const userID = int64(1001) + tests := []struct { + name string + media tg.InputMediaClass + }{ + { + name: "uploaded photo", + media: &tg.InputMediaUploadedPhoto{File: &tg.InputFile{ID: 11, Parts: 1, Name: "gone.jpg"}}, + }, + { + name: "referenced photo", + media: &tg.InputMediaPhoto{ID: &tg.InputPhoto{ID: 22, AccessHash: 220}}, + }, + { + name: "poll", + media: &tg.InputMediaPoll{Poll: tg.Poll{ + Question: tg.TextWithEntities{Text: "question?", Entities: []tg.MessageEntityClass{}}, + Answers: []tg.PollAnswerClass{ + &tg.PollAnswer{Text: tg.TextWithEntities{Text: "yes", Entities: []tg.MessageEntityClass{}}, Option: []byte{0}}, + &tg.PollAnswer{Text: tg.TextWithEntities{Text: "no", Entities: []tg.MessageEntityClass{}}, Option: []byte{1}}, + }, + }}, + }, + } + for i, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + randomID := int64(7100 + i) + messages := newReplayPreflightMessages() + messages.replays[randomID] = privateReplayFixture(userID, userID, randomID, 80+i, 0) + limiter := &captureRateLimiter{block: true, retryAfter: 30} + r := New(Config{SendRateLimit: 1, SendRateWindow: time.Minute}, Deps{ + Messages: messages, + Limiter: limiter, + // Files and Polls are intentionally nil. Reaching any resolver would fail. + }, zaptest.NewLogger(t), clock.System) + req := &tg.MessagesSendMediaRequest{ + Peer: &tg.InputPeerSelf{}, + Media: tc.media, + Message: "already committed", + RandomID: randomID, + } + if fingerprint, err := sendMediaIdempotencyFingerprint(req); err != nil || len(fingerprint) != 32 { + t.Fatalf("fingerprint len=%d err=%v, want SHA-256", len(fingerprint), err) + } + if _, err := r.onMessagesSendMedia(WithUserID(context.Background(), userID), req); err != nil { + t.Fatalf("exact replay: %v", err) + } + if len(limiter.calls) != 0 { + t.Fatalf("limiter calls = %+v, want none", limiter.calls) + } + if len(messages.sendRequests) != 0 || len(messages.lookupRequests) != 1 { + t.Fatalf("lookup=%d send=%d, want 1/0", len(messages.lookupRequests), len(messages.sendRequests)) + } + }) + } +} + +func TestSendMultiMediaMixedReplayChargesAndResolvesOnlyAbsentItems(t *testing.T) { + const userID = int64(1001) + messages := newReplayPreflightMessages() + messages.replays[7201] = privateReplayFixture(userID, userID, 7201, 91, messages.reservedGroupedID) + limiter := &captureRateLimiter{} + dialogs := &replayCountingDialogs{captureDialogs: &captureDialogs{}} + files := &replaySelectiveFiles{fakeFiles: &fakeFiles{photos: map[int64]domain.Photo{ + 222: {ID: 222, AccessHash: 2220, DCID: 2, Sizes: []domain.PhotoSize{{Kind: domain.PhotoSizeKindDefault, Type: "x", W: 100, H: 100}}}, + }}} + r := New(Config{SendRateLimit: 10, SendRateWindow: time.Minute}, Deps{ + Messages: messages, + Dialogs: dialogs, + Files: files, + Limiter: limiter, + }, zaptest.NewLogger(t), clock.System) + req := &tg.MessagesSendMultiMediaRequest{ + Peer: &tg.InputPeerSelf{}, + ClearDraft: true, + MultiMedia: []tg.InputSingleMedia{ + {Media: &tg.InputMediaPhoto{ID: &tg.InputPhoto{ID: 111, AccessHash: 1110}}, RandomID: 7201, Message: "duplicate"}, + {Media: &tg.InputMediaPhoto{ID: &tg.InputPhoto{ID: 222, AccessHash: 2220}}, RandomID: 7202, Message: "new"}, + }, + } + + if _, err := r.onMessagesSendMultiMedia(WithUserID(context.Background(), userID), req); err != nil { + t.Fatalf("mixed sendMultiMedia: %v", err) + } + if len(limiter.calls) != 1 || limiter.calls[0].cost != 1 { + t.Fatalf("limiter calls = %+v, want one absent-item cost", limiter.calls) + } + if !reflect.DeepEqual(files.getPhotoIDs, []int64{222}) { + t.Fatalf("resolved photo IDs = %v, want only absent item 222", files.getPhotoIDs) + } + if len(messages.sendRequests) != 1 || messages.sendRequests[0].RandomID != 7202 { + t.Fatalf("send requests = %+v, want only random_id 7202", messages.sendRequests) + } + if len(messages.reserveRequests) != 1 { + t.Fatalf("album reservations = %d, want 1", len(messages.reserveRequests)) + } + if dialogs.deleteDraftCalls != 1 { + t.Fatalf("draft deletes = %d, want exactly one from first genuinely-new item", dialogs.deleteDraftCalls) + } + + limiterCalls := len(limiter.calls) + resolved := append([]int64(nil), files.getPhotoIDs...) + reservations := len(messages.reserveRequests) + if _, err := r.onMessagesSendMultiMedia(WithUserID(context.Background(), userID), req); err != nil { + t.Fatalf("full duplicate sendMultiMedia: %v", err) + } + if len(limiter.calls) != limiterCalls { + t.Fatalf("full duplicate added limiter calls: before=%d after=%d", limiterCalls, len(limiter.calls)) + } + if !reflect.DeepEqual(files.getPhotoIDs, resolved) { + t.Fatalf("full duplicate resolved media: before=%v after=%v", resolved, files.getPhotoIDs) + } + if len(messages.reserveRequests) != reservations { + t.Fatalf("full duplicate reservations: before=%d after=%d", reservations, len(messages.reserveRequests)) + } + if dialogs.deleteDraftCalls != 1 { + t.Fatalf("full duplicate cleared draft again: calls=%d", dialogs.deleteDraftCalls) + } +} + +func TestForwardReplayPreflightSkipsCommittedSourcesAndLoadsOnlyAbsentIDs(t *testing.T) { + const userID = int64(1001) + t.Run("full duplicate tolerates deleted sources", func(t *testing.T) { + messages := newReplayPreflightMessages() + messages.replays[7301] = privateReplayFixture(userID, userID, 7301, 101, 0) + messages.replays[7302] = privateReplayFixture(userID, userID, 7302, 102, 0) + limiter := &captureRateLimiter{block: true, retryAfter: 30} + r := New(Config{SendRateLimit: 1, SendRateWindow: time.Minute}, Deps{Messages: messages, Limiter: limiter}, zaptest.NewLogger(t), clock.System) + + if _, err := r.onMessagesForwardMessages(WithUserID(context.Background(), userID), &tg.MessagesForwardMessagesRequest{ + FromPeer: &tg.InputPeerEmpty{}, + ToPeer: &tg.InputPeerSelf{}, + ID: []int{41, 42}, + RandomID: []int64{7301, 7302}, + }); err != nil { + t.Fatalf("full duplicate forward with deleted sources: %v", err) + } + if messages.getMessagesCalls != 0 { + t.Fatalf("GetMessages calls = %d, want 0", messages.getMessagesCalls) + } + if len(limiter.calls) != 0 || len(messages.sendRequests) != 0 { + t.Fatalf("limiter=%v sends=%d, want no side effects", limiter.calls, len(messages.sendRequests)) + } + }) + + t.Run("mixed loads absent IDs only", func(t *testing.T) { + messages := newReplayPreflightMessages() + messages.replays[7311] = privateReplayFixture(userID, userID, 7311, 111, 0) + messages.list = domain.MessageList{Messages: []domain.Message{{ + ID: 52, + OwnerUserID: userID, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: 2002}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: 2002}, + Date: 1700000100, + Body: "still available", + }}} + limiter := &captureRateLimiter{} + r := New(Config{SendRateLimit: 10, SendRateWindow: time.Minute}, Deps{Messages: messages, Limiter: limiter}, zaptest.NewLogger(t), clock.System) + + if _, err := r.onMessagesForwardMessages(WithUserID(context.Background(), userID), &tg.MessagesForwardMessagesRequest{ + FromPeer: &tg.InputPeerEmpty{}, + ToPeer: &tg.InputPeerSelf{}, + ID: []int{51, 52}, + RandomID: []int64{7311, 7312}, + }); err != nil { + t.Fatalf("mixed forward: %v", err) + } + if messages.getMessagesCalls != 1 || !reflect.DeepEqual(messages.getMessagesIDs, [][]int{{52}}) { + t.Fatalf("GetMessages calls=%d ids=%v, want one [52]", messages.getMessagesCalls, messages.getMessagesIDs) + } + if len(limiter.calls) != 1 || limiter.calls[0].cost != 1 { + t.Fatalf("limiter calls = %+v, want cost 1", limiter.calls) + } + if len(messages.sendRequests) != 1 || messages.sendRequests[0].RandomID != 7312 { + t.Fatalf("send requests = %+v, want only random_id 7312", messages.sendRequests) + } + }) +} + +func TestChannelAndMonoforumExactReplayPrecedesLimiterAndCurrentPermission(t *testing.T) { + t.Run("channel replay survives sender ban", func(t *testing.T) { + ctx := context.Background() + users := memory.NewUserStore() + owner, err := users.Create(ctx, domain.User{AccessHash: 11, Phone: "15550007001", FirstName: "Owner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + member, err := users.Create(ctx, domain.User{AccessHash: 12, Phone: "15550007002", FirstName: "Member"}) + if err != nil { + t.Fatalf("create member: %v", err) + } + channels := appchannels.NewService(memory.NewChannelStore()) + created, err := channels.CreateMegagroupFromCreateChat(ctx, owner.ID, domain.CreateChannelRequest{ + Title: "Replay Group", MemberUserIDs: []int64{member.ID}, Date: 100, + }) + if err != nil { + t.Fatalf("create group: %v", err) + } + limiter := &captureRateLimiter{} + r := New(Config{SendRateLimit: 10, SendRateWindow: time.Minute}, Deps{ + Users: appusers.NewService(users), + Channels: channels, + Limiter: limiter, + }, zaptest.NewLogger(t), clock.System) + req := &tg.MessagesSendMessageRequest{ + Peer: &tg.InputPeerChannel{ChannelID: created.Channel.ID, AccessHash: created.Channel.AccessHash}, + Message: "committed before ban", + RandomID: 7401, + } + memberCtx := WithUserID(ctx, member.ID) + if _, err := r.onMessagesSendMessage(memberCtx, req); err != nil { + t.Fatalf("first channel send: %v", err) + } + callsBeforeReplay := len(limiter.calls) + if _, err := channels.EditBanned(ctx, owner.ID, domain.EditChannelBannedRequest{ + ChannelID: created.Channel.ID, + Participant: domain.Peer{Type: domain.PeerTypeUser, ID: member.ID}, + BannedRights: domain.ChannelBannedRights{ + ViewMessages: true, + UntilDate: 1000, + }, + Date: 101, + }); err != nil { + t.Fatalf("ban member: %v", err) + } + limiter.block = true + if _, err := r.onMessagesSendMessage(memberCtx, req); err != nil { + t.Fatalf("channel exact replay after ban: %v", err) + } + if len(limiter.calls) != callsBeforeReplay { + t.Fatalf("replay limiter calls: before=%d after=%d", callsBeforeReplay, len(limiter.calls)) + } + }) + + t.Run("monoforum replay survives direct-message disable", func(t *testing.T) { + ctx := context.Background() + users := memory.NewUserStore() + owner, err := users.Create(ctx, domain.User{AccessHash: 21, Phone: "15550007101", FirstName: "Owner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + subscriber, err := users.Create(ctx, domain.User{AccessHash: 22, Phone: "15550007102", FirstName: "Subscriber"}) + if err != nil { + t.Fatalf("create subscriber: %v", err) + } + channelStore := memory.NewChannelStore() + channels := appchannels.NewService(channelStore) + created, err := channels.CreateChannel(ctx, owner.ID, domain.CreateChannelRequest{Title: "Replay DM", Broadcast: true, Date: 200}) + if err != nil { + t.Fatalf("create broadcast: %v", err) + } + enabled, err := channelStore.SetPaidMessagesPrice(ctx, owner.ID, created.Channel.ID, 0, true) + if err != nil { + t.Fatalf("enable direct messages: %v", err) + } + monoID := enabled.Channel.LinkedMonoforumID + limiter := &captureRateLimiter{} + r := New(Config{SendRateLimit: 10, SendRateWindow: time.Minute}, Deps{ + Users: appusers.NewService(users), + Channels: channels, + Limiter: limiter, + }, zaptest.NewLogger(t), clock.System) + req := &tg.MessagesSendMessageRequest{ + Peer: &tg.InputPeerChannel{ChannelID: monoID}, + Message: "committed before disable", + RandomID: 7411, + } + req.SetReplyTo(&tg.InputReplyToMonoForum{MonoforumPeerID: &tg.InputPeerUser{UserID: subscriber.ID}}) + subscriberCtx := WithUserID(ctx, subscriber.ID) + if _, err := r.onMessagesSendMessage(subscriberCtx, req); err != nil { + t.Fatalf("first monoforum send: %v", err) + } + callsBeforeReplay := len(limiter.calls) + if _, err := channelStore.SetPaidMessagesPrice(ctx, owner.ID, created.Channel.ID, 0, false); err != nil { + t.Fatalf("disable direct messages: %v", err) + } + limiter.block = true + if _, err := r.onMessagesSendMessage(subscriberCtx, req); err != nil { + t.Fatalf("monoforum exact replay after disable: %v", err) + } + if len(limiter.calls) != callsBeforeReplay { + t.Fatalf("replay limiter calls: before=%d after=%d", callsBeforeReplay, len(limiter.calls)) + } + }) +} diff --git a/internal/rpc/messages_bot_longtail.go b/internal/rpc/messages_bot_longtail.go index de592f2f..c75f2f52 100644 --- a/internal/rpc/messages_bot_longtail.go +++ b/internal/rpc/messages_bot_longtail.go @@ -39,12 +39,16 @@ func (r *Router) onMessagesSendWebViewData(ctx context.Context, req *tg.Messages if err != nil { return nil, err } - authKeyID, _ := AuthKeyIDFrom(ctx) + idempotencyFingerprint, err := rpcRequestFingerprint(req) + if err != nil { + return nil, internalErr() + } sessionID, _ := SessionIDFrom(ctx) res, err := r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{ - SenderUserID: userID, - RecipientUserID: bot.ID, - RandomID: req.RandomID, + SenderUserID: userID, + RecipientUserID: bot.ID, + RandomID: req.RandomID, + IdempotencyFingerprint: idempotencyFingerprint, Media: &domain.MessageMedia{ Kind: domain.MessageMediaKindService, ServiceAction: &domain.MessageServiceAction{ @@ -56,16 +60,20 @@ func (r *Router) onMessagesSendWebViewData(ctx context.Context, req *tg.Messages }, }, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) if err != nil { - return nil, internalErr() + return nil, messageSendErr(err) } - users := r.usersForMessageUpdate(ctx, userID, res.SenderMessage) - chats := r.chatsForMessageUpdate(ctx, userID, res.SenderMessage) - return tgPrivateMessageUpdates(res.SenderEvent, res.SenderMessage, 0, false, users, chats), nil + var users []tg.UserClass + var chats []tg.ChatClass + if !res.Duplicate { + users = r.usersForMessageUpdate(ctx, userID, res.SenderMessage) + chats = r.chatsForMessageUpdate(ctx, userID, res.SenderMessage) + } + return tgPrivateSendResultUpdates(res, req.RandomID, false, users, chats), nil } func (r *Router) onMessagesSendBotRequestedPeer(ctx context.Context, req *tg.MessagesSendBotRequestedPeerRequest) (tg.UpdatesClass, error) { @@ -122,7 +130,6 @@ func (r *Router) onMessagesSendBotRequestedPeer(ctx context.Context, req *tg.Mes if err != nil { return nil, err } - authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) res, err := r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{ SenderUserID: userID, @@ -139,7 +146,7 @@ func (r *Router) onMessagesSendBotRequestedPeer(ctx context.Context, req *tg.Mes }, }, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) @@ -147,9 +154,13 @@ func (r *Router) onMessagesSendBotRequestedPeer(ctx context.Context, req *tg.Mes return nil, internalErr() } _ = r.deps.Bots.DeleteRequestedWebViewButton(ctx, botUser.ID, userID, webAppReqID) - users := r.usersForMessageUpdate(ctx, userID, res.SenderMessage) - chats := r.chatsForMessageUpdate(ctx, userID, res.SenderMessage) - return tgPrivateMessageUpdates(res.SenderEvent, res.SenderMessage, 0, false, users, chats), nil + var users []tg.UserClass + var chats []tg.ChatClass + if !res.Duplicate { + users = r.usersForMessageUpdate(ctx, userID, res.SenderMessage) + chats = r.chatsForMessageUpdate(ctx, userID, res.SenderMessage) + } + return tgPrivateSendResultUpdates(res, res.SenderMessage.RandomID, false, users, chats), nil } func requestedPeerTypeMatches(kind string, peer domain.Peer) bool { diff --git a/internal/rpc/messages_bot_longtail_rpc_test.go b/internal/rpc/messages_bot_longtail_rpc_test.go index 996879f0..e4386f88 100644 --- a/internal/rpc/messages_bot_longtail_rpc_test.go +++ b/internal/rpc/messages_bot_longtail_rpc_test.go @@ -67,15 +67,22 @@ func TestMessagesSendWebViewDataServiceMessageRoundTrip(t *testing.T) { repeat, err := f.router.onMessagesSendWebViewData(ownerCtx, &tg.MessagesSendWebViewDataRequest{ Bot: inputUser(f.bot), RandomID: 1001, - ButtonText: "Open changed", - Data: `{"ok":false}`, + ButtonText: "Open", + Data: `{"ok":true}`, }) if err != nil { t.Fatalf("repeat send webview data: %v", err) } - repeatMsg := repeat.(*tg.Updates).Updates[0].(*tg.UpdateNewMessage).Message.(*tg.MessageService) - if repeatMsg.ID != service.ID { - t.Fatalf("repeat message id = %d, want original %d", repeatMsg.ID, service.ID) + repeatUpdates := repeat.(*tg.Updates).Updates + if len(repeatUpdates) != 2 { + t.Fatalf("repeat updates len = %d, want immutable mapping + message confirmation", len(repeatUpdates)) + } + mapping, ok := repeatUpdates[0].(*tg.UpdateMessageID) + if !ok || mapping.ID != service.ID || mapping.RandomID != 1001 { + t.Fatalf("repeat mapping = %+v (%T), want random_id 1001 -> original %d", repeatUpdates[0], repeatUpdates[0], service.ID) + } + if replayed, ok := repeatUpdates[1].(*tg.UpdateNewMessage); !ok || replayed.Pts <= 0 || replayed.PtsCount != 1 { + t.Fatalf("repeat confirmation = %T %+v, want UpdateNewMessage", repeatUpdates[1], repeatUpdates[1]) } botHistory, err = f.router.deps.Messages.GetHistory(ctx, f.bot.ID, domain.MessageFilter{ HasPeer: true, @@ -85,6 +92,14 @@ func TestMessagesSendWebViewDataServiceMessageRoundTrip(t *testing.T) { if err != nil || len(botHistory.Messages) != 1 { t.Fatalf("bot history after repeat = %+v err=%v, want still one service message", botHistory, err) } + if _, err := f.router.onMessagesSendWebViewData(ownerCtx, &tg.MessagesSendWebViewDataRequest{ + Bot: inputUser(f.bot), + RandomID: 1001, + ButtonText: "Open changed", + Data: `{"ok":false}`, + }); !tgerr.Is(err, "RANDOM_ID_DUPLICATE") { + t.Fatalf("conflicting webview random_id err = %v, want RANDOM_ID_DUPLICATE", err) + } if _, err := f.router.onMessagesSendWebViewData(ownerCtx, &tg.MessagesSendWebViewDataRequest{ Bot: inputUser(f.peer), diff --git a/internal/rpc/messages_compat.go b/internal/rpc/messages_compat.go index a36f6f43..bf8a776d 100644 --- a/internal/rpc/messages_compat.go +++ b/internal/rpc/messages_compat.go @@ -140,7 +140,7 @@ func (r *Router) onMessagesDeleteSavedHistory(ctx context.Context, req *tg.Messa MinDate: minDate, MaxDate: maxDate, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { diff --git a/internal/rpc/messages_delete.go b/internal/rpc/messages_delete.go index d0fa0c37..adea5ce9 100644 --- a/internal/rpc/messages_delete.go +++ b/internal/rpc/messages_delete.go @@ -27,7 +27,7 @@ func (r *Router) onMessagesDeleteMessages(ctx context.Context, req *tg.MessagesD IDs: req.ID, Revoke: req.GetRevoke(), Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { @@ -102,7 +102,7 @@ func (r *Router) onMessagesDeleteHistory(ctx context.Context, req *tg.MessagesDe JustClear: req.GetJustClear(), Revoke: req.GetRevoke(), Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { diff --git a/internal/rpc/messages_dialogs.go b/internal/rpc/messages_dialogs.go index 9c962a3b..ade702fe 100644 --- a/internal/rpc/messages_dialogs.go +++ b/internal/rpc/messages_dialogs.go @@ -68,7 +68,7 @@ func (r *Router) recordDraftMessageEvent(ctx context.Context, userID int64, peer } authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordDraftMessage(ctx, authKeyID, userID, peer, topMsgID, sessionID) + event, state, err := r.deps.Updates.RecordDraftMessage(ctx, authKeyID, userID, peer, topMsgID, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { r.log.Warn("record draft message event", zap.Int64("user_id", userID), zap.Error(err)) return domain.UpdateEvent{} @@ -433,7 +433,7 @@ func (r *Router) onMessagesUpdateDialogFilter(ctx context.Context, req *tg.Messa if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, req.ID, folder, sessionID) + event, _, err = r.deps.Updates.RecordDialogFilter(ctx, authKeyID, userID, req.ID, folder, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -465,7 +465,7 @@ func (r *Router) onMessagesUpdateDialogFiltersOrder(ctx context.Context, order [ if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err = r.deps.Updates.RecordDialogFilterOrder(ctx, authKeyID, userID, clean, sessionID) + event, _, err = r.deps.Updates.RecordDialogFilterOrder(ctx, authKeyID, userID, clean, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -493,7 +493,7 @@ func (r *Router) onMessagesToggleDialogFilterTags(ctx context.Context, enabled b if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err = r.deps.Updates.RecordDialogFiltersReload(ctx, authKeyID, userID, sessionID) + event, _, err = r.deps.Updates.RecordDialogFiltersReload(ctx, authKeyID, userID, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -584,7 +584,7 @@ func (r *Router) onMessagesToggleDialogPin(ctx context.Context, req *tg.Messages if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordDialogPinned(ctx, authKeyID, userID, peers[0], pinned, folderID, sessionID) + event, state, err := r.deps.Updates.RecordDialogPinned(ctx, authKeyID, userID, peers[0], pinned, folderID, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -631,7 +631,7 @@ func (r *Router) toggleArchiveFolderPin(ctx context.Context, userID int64, folde if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordDialogPinned(ctx, authKeyID, userID, folderPeer, pinned, 0, sessionID) + event, state, err := r.deps.Updates.RecordDialogPinned(ctx, authKeyID, userID, folderPeer, pinned, 0, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -682,7 +682,7 @@ func (r *Router) onMessagesReorderPinnedDialogs(ctx context.Context, req *tg.Mes if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordPinnedDialogs(ctx, authKeyID, userID, req.FolderID, peers, sessionID) + event, state, err := r.deps.Updates.RecordPinnedDialogs(ctx, authKeyID, userID, req.FolderID, peers, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -741,7 +741,7 @@ func (r *Router) onMessagesMarkDialogUnread(ctx context.Context, req *tg.Message if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordDialogUnreadMark(ctx, authKeyID, userID, peers[0], unread, sessionID) + event, state, err := r.deps.Updates.RecordDialogUnreadMark(ctx, authKeyID, userID, peers[0], unread, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -819,7 +819,7 @@ func (r *Router) onMessagesHidePeerSettingsBar(ctx context.Context, input tg.Inp if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, domain.PeerSettings{HiddenPeerSettingsBar: true}, sessionID) + event, state, err := r.deps.Updates.RecordPeerSettings(ctx, authKeyID, userID, peer, domain.PeerSettings{HiddenPeerSettingsBar: true}, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } diff --git a/internal/rpc/messages_edit.go b/internal/rpc/messages_edit.go index ad49ff59..28d15e1c 100644 --- a/internal/rpc/messages_edit.go +++ b/internal/rpc/messages_edit.go @@ -143,7 +143,6 @@ func (r *Router) onMessagesEditMessage(ctx context.Context, req *tg.MessagesEdit return nil, messageEditForbiddenErr() } sessionID, _ := SessionIDFrom(ctx) - authKeyID, _ := AuthKeyIDFrom(ctx) res, err := r.deps.Messages.EditMessage(ctx, userID, domain.EditMessageRequest{ OwnerUserID: userID, Peer: peer, @@ -151,7 +150,7 @@ func (r *Router) onMessagesEditMessage(ctx context.Context, req *tg.MessagesEdit Message: message, Entities: domainMessageEntitiesForViewer(userID, entities), EditDate: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, SetReplyMarkup: setReplyMarkup, ReplyMarkup: replyMarkup, diff --git a/internal/rpc/messages_edit_geolive.go b/internal/rpc/messages_edit_geolive.go index 95751ece..a09c28c0 100644 --- a/internal/rpc/messages_edit_geolive.go +++ b/internal/rpc/messages_edit_geolive.go @@ -94,7 +94,6 @@ func (r *Router) onEditMessageLiveLocation(ctx context.Context, req *tg.Messages return nil, peerIDInvalidErr() } sessionID, _ := SessionIDFrom(ctx) - authKeyID, _ := AuthKeyIDFrom(ctx) res, err := r.deps.Messages.EditMessage(ctx, userID, domain.EditMessageRequest{ OwnerUserID: userID, Peer: peer, @@ -103,7 +102,7 @@ func (r *Router) onEditMessageLiveLocation(ctx context.Context, req *tg.Messages Entities: current.entities, Media: newMedia, EditDate: now, - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { diff --git a/internal/rpc/messages_forum_read.go b/internal/rpc/messages_forum_read.go index 12181f73..16dbda61 100644 --- a/internal/rpc/messages_forum_read.go +++ b/internal/rpc/messages_forum_read.go @@ -61,7 +61,7 @@ func (r *Router) recordChannelDiscussionInbox(ctx context.Context, userID, chann if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - rec, _, err := r.deps.Updates.RecordChannelDiscussionInbox(ctx, authKeyID, userID, channelID, topicID, maxID, sessionID) + rec, _, err := r.deps.Updates.RecordChannelDiscussionInbox(ctx, authKeyID, userID, channelID, topicID, maxID, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return internalErr() } diff --git a/internal/rpc/messages_forward.go b/internal/rpc/messages_forward.go index fccfccde..4ac9a8e6 100644 --- a/internal/rpc/messages_forward.go +++ b/internal/rpc/messages_forward.go @@ -43,11 +43,80 @@ func (r *Router) onMessagesForwardMessages(ctx context.Context, req *tg.Messages if userID == 0 { return nil, peerIDInvalidErr() } - fromPeer, preloadedSources, err := r.forwardFromPeerAndSources(ctx, userID, req.FromPeer, req.ID, req.RandomID) + if !forwardMessageIDsValid(req.ID, req.RandomID) { + return nil, messageIDInvalidErr() + } + seenRandomIDs := make(map[int64]struct{}, len(req.RandomID)) + for _, randomID := range req.RandomID { + if _, duplicate := seenRandomIDs[randomID]; duplicate { + return nil, randomIDDuplicateErr() + } + seenRandomIDs[randomID] = struct{}{} + } + toPeer, ok := r.domainPeerFromInputPeer(userID, req.ToPeer) + if !ok || toPeer.ID == 0 { + return nil, peerIDInvalidErr() + } + idempotencyFingerprints := make([][]byte, len(req.ID)) + for i := range req.ID { + idempotencyFingerprints[i], err = forwardMessagesItemIdempotencyFingerprint(req, req.ID[i], req.RandomID[i]) + if err != nil { + return nil, internalErr() + } + } + immediate := req.ScheduleDate == 0 || scheduleDateIsImmediate(req.ScheduleDate, int(r.clock.Now().Unix())) + replays := make([]outgoingReplayLookup, len(req.ID)) + absentIndexes := make([]int, 0, len(req.ID)) + if immediate { + for i := range req.ID { + replay, err := r.lookupOutgoingReplay(ctx, userID, toPeer, req.RandomID[i], idempotencyFingerprints[i]) + if err != nil { + return nil, err + } + replays[i] = replay + if !replay.found { + absentIndexes = append(absentIndexes, i) + } + } + } else { + for i := range req.ID { + absentIndexes = append(absentIndexes, i) + } + } + if len(absentIndexes) == 0 { + if toPeer.Type == domain.PeerTypeChannel { + results := make([]domain.SendChannelMessageResult, len(replays)) + for i := range replays { + results[i] = replays[i].channel + } + return r.channelMessagesUpdatesWithPeerCache(ctx, userID, results, req.RandomID, true, nil, newViewerPeerCache(r)), nil + } + res := domain.ForwardPrivateMessagesResult{OwnerUserID: userID} + for i := range replays { + sent := replays[i].private + res.SenderMessages = append(res.SenderMessages, sent.SenderMessage) + res.RecipientMessages = append(res.RecipientMessages, sent.RecipientMessage) + res.SenderEvents = append(res.SenderEvents, sent.SenderEvent) + res.RecipientEvents = append(res.RecipientEvents, sent.RecipientEvent) + res.Duplicates = append(res.Duplicates, true) + res.ReplayDeleteEvents = append(res.ReplayDeleteEvents, sent.ReplayDeleteEvent) + } + return tgForwardMessagesUpdates(res, req.RandomID, r.usersForMessageUpdates(ctx, userID, res.SenderMessages), r.chatsForMessageUpdates(ctx, userID, res.SenderMessages)), nil + } + if err := r.checkSendRateLimit(ctx, userID, len(absentIndexes)); err != nil { + return nil, err + } + toPeer, err = r.checkedDomainPeerFromInputPeer(ctx, userID, req.ToPeer) if err != nil { return nil, err } - toPeer, err := r.checkedDomainPeerFromInputPeer(ctx, userID, req.ToPeer) + absentIDs := make([]int, len(absentIndexes)) + absentRandomIDs := make([]int64, len(absentIndexes)) + for i, originalIndex := range absentIndexes { + absentIDs[i] = req.ID[originalIndex] + absentRandomIDs[i] = req.RandomID[originalIndex] + } + fromPeer, preloadedSources, err := r.forwardFromPeerAndSources(ctx, userID, req.FromPeer, absentIDs, absentRandomIDs) if err != nil { return nil, err } @@ -75,72 +144,89 @@ func (r *Router) onMessagesForwardMessages(ctx context.Context, req *tg.Messages } } } - if !forwardMessageIDsValid(req.ID, req.RandomID) { - return nil, messageIDInvalidErr() - } - if err := r.checkSendRateLimit(ctx, userID, len(req.ID)); err != nil { - return nil, err - } if req.ScheduleDate != 0 && !scheduleDateIsImmediate(req.ScheduleDate, int(r.clock.Now().Unix())) { return r.scheduleForwardMessages(ctx, userID, fromPeer, toPeer, req, replyTo, sendAs, preloadedSources) } + absentSources, err := r.forwardSourcesForRequest(ctx, userID, fromPeer, absentIDs, preloadedSources) + if err != nil { + return nil, messageForwardErr(err) + } + sources := make([]forwardSource, len(req.ID)) + for i, originalIndex := range absentIndexes { + sources[originalIndex] = absentSources[i] + } if toPeer.Type == domain.PeerTypeChannel { if r.deps.Channels == nil { return nil, peerIDInvalidErr() } - sources, err := r.forwardSourcesForRequest(ctx, userID, fromPeer, req.ID, preloadedSources) - if err != nil { - return nil, messageForwardErr(err) - } recipients := make([]int64, 0) results := make([]domain.SendChannelMessageResult, 0, len(sources)) extraUserIDs := make([]int64, 0, len(sources)) + fanoutResults := make([]domain.SendChannelMessageResult, 0, len(sources)) + fanoutExtraUserIDs := make([]int64, 0, len(sources)) for i, source := range sources { + if replays[i].found { + results = append(results, replays[i].channel) + continue + } forward := source.forward if req.DropAuthor { forward = nil } mentionUserIDs := r.mentionUserIDsFromDomain(ctx, userID, source.body, source.entities) res, err := r.deps.Channels.SendMessage(ctx, userID, domain.SendChannelMessageRequest{ - UserID: userID, - ChannelID: toPeer.ID, - RandomID: req.RandomID[i], - Message: source.body, - Entities: source.entities, - Media: source.media, - MentionUserIDs: mentionUserIDs, - Silent: req.Silent, - NoForwards: req.Noforwards, - ReplyTo: replyTo, - Forward: forward, - SendAs: sendAs, - Date: int(r.clock.Now().Unix()), + UserID: userID, + ChannelID: toPeer.ID, + RandomID: req.RandomID[i], + IdempotencyFingerprint: idempotencyFingerprints[i], + IdempotencyPreflighted: replays[i].checked, + Message: source.body, + Entities: source.entities, + Media: source.media, + MentionUserIDs: mentionUserIDs, + Silent: req.Silent, + NoForwards: req.Noforwards, + ReplyTo: replyTo, + Forward: forward, + SendAs: sendAs, + Date: int(r.clock.Now().Unix()), }) if err != nil { return nil, channelInvalidErr(err) } results = append(results, res) + sourceUserID := source.userID() + if sourceUserID != 0 { + extraUserIDs = append(extraUserIDs, sourceUserID) + } + // An exact random_id replay must still appear in the caller's echo, but it must not + // emit the old channel pts as a fresh realtime payload/Bot API update/discussion push. + // Mixed batches therefore fan out only newly committed results while preserving the + // complete result vector for TDesktop's random_id -> message-id reconciliation. + if res.Duplicate { + continue + } + fanoutResults = append(fanoutResults, res) + if sourceUserID != 0 { + fanoutExtraUserIDs = append(fanoutExtraUserIDs, sourceUserID) + } // 收件人是该频道的活跃成员集，对本次转发的每条源都相同；只取一次， // 避免一次转发 ≤100 条到 N 成员大群时把 recipients 累积成 ~100×N 条目 // 的巨大临时切片（N=10^5 时约千万级）。 if len(recipients) == 0 { recipients = res.Recipients } - if sourceUserID := source.userID(); sourceUserID != 0 { - extraUserIDs = append(extraUserIDs, sourceUserID) - } } // echo 与 fan-out 用各自独立 cache（RPC vs worker goroutine 防竞态）。多条转发汇成 // 一个 fan-out job（channelMessagesUpdatesWithPeerCache 内含多条 UpdateNewChannelMessage）， // 由同 channel 分片 FIFO 原子投递。 echoCache := newViewerPeerCache(r) updates := r.channelMessagesUpdatesWithPeerCache(ctx, userID, results, req.RandomID, true, extraUserIDs, echoCache) - fanoutPts := 0 - if n := len(results); n > 0 { - fanoutPts = results[n-1].Event.Pts + if n := len(fanoutResults); n > 0 { + fanoutPts := fanoutResults[n-1].Event.Pts + r.enqueueChannelMessagesFanout(ctx, userID, toPeer.ID, fanoutPts, recipients, fanoutResults, fanoutExtraUserIDs) } - r.enqueueChannelMessagesFanout(ctx, userID, toPeer.ID, fanoutPts, recipients, results, extraUserIDs) - for _, res := range results { + for _, res := range fanoutResults { r.pushChannelDiscussionUpdate(ctx, userID, res.Discussion) } return updates, nil @@ -153,17 +239,20 @@ func (r *Router) onMessagesForwardMessages(ctx context.Context, req *tg.Messages if err != nil { return nil, err } - // 私聊源与频道源统一经 forwardSources 取源：首次生成的 forward header 在 - // forwardSources 内已按原作者 PrivacyKeyForwards 降级（不允许链接回账号时仅保 - // 留 from_name），避免私聊→私聊路径泄漏原作者可点击账号；media 也随 source 透传。 - sources, err := r.forwardSourcesForRequest(ctx, userID, fromPeer, req.ID, preloadedSources) - if err != nil { - return nil, messageForwardErr(err) - } sessionID, _ := SessionIDFrom(ctx) - authKeyID, _ := AuthKeyIDFrom(ctx) + authKeyID := rawAuthKeyIDForOrigin(ctx) res := domain.ForwardPrivateMessagesResult{OwnerUserID: userID} for i, source := range sources { + if replays[i].found { + sent := replays[i].private + res.SenderMessages = append(res.SenderMessages, sent.SenderMessage) + res.RecipientMessages = append(res.RecipientMessages, sent.RecipientMessage) + res.SenderEvents = append(res.SenderEvents, sent.SenderEvent) + res.RecipientEvents = append(res.RecipientEvents, sent.RecipientEvent) + res.Duplicates = append(res.Duplicates, true) + res.ReplayDeleteEvents = append(res.ReplayDeleteEvents, sent.ReplayDeleteEvent) + continue + } forward := source.forward if req.DropAuthor { forward = nil @@ -175,20 +264,22 @@ func (r *Router) onMessagesForwardMessages(ctx context.Context, req *tg.Messages forward = &saved } sent, err := r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{ - SenderUserID: userID, - RecipientUserID: toPeer.ID, - RandomID: req.RandomID[i], - Message: source.body, - Entities: source.entities, - Media: source.media, - Silent: req.Silent, - NoForwards: req.Noforwards, - ReplyTo: replyTo, - Forward: forward, - Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, - OriginSessionID: sessionID, - RecipientBlocked: recipientBlocked, + SenderUserID: userID, + RecipientUserID: toPeer.ID, + RandomID: req.RandomID[i], + Message: source.body, + Entities: source.entities, + Media: source.media, + Silent: req.Silent, + NoForwards: req.Noforwards, + ReplyTo: replyTo, + Forward: forward, + Date: int(r.clock.Now().Unix()), + OriginAuthKeyID: authKeyID, + OriginSessionID: sessionID, + RecipientBlocked: recipientBlocked, + IdempotencyFingerprint: idempotencyFingerprints[i], + IdempotencyPreflighted: replays[i].checked, }) if err != nil { return nil, messageForwardErr(err) @@ -201,6 +292,7 @@ func (r *Router) onMessagesForwardMessages(ctx context.Context, req *tg.Messages res.SenderEvents = append(res.SenderEvents, sent.SenderEvent) res.RecipientEvents = append(res.RecipientEvents, sent.RecipientEvent) res.Duplicates = append(res.Duplicates, sent.Duplicate) + res.ReplayDeleteEvents = append(res.ReplayDeleteEvents, sent.ReplayDeleteEvent) } return tgForwardMessagesUpdates(res, req.RandomID, r.usersForMessageUpdates(ctx, userID, res.SenderMessages), r.chatsForMessageUpdates(ctx, userID, res.SenderMessages)), nil } @@ -497,6 +589,8 @@ func messageForwardErr(err error) error { return chatForwardsRestrictedErr() case errors.Is(err, domain.ErrReplyMessageIDInvalid): return replyMessageIDInvalidErr() + case errors.Is(err, domain.ErrMessageRandomIDDuplicate): + return randomIDDuplicateErr() default: return internalErr() } @@ -532,6 +626,18 @@ func tgForwardMessagesUpdates(res domain.ForwardPrivateMessagesResult, randomIDs Pts: pts, PtsCount: ptsCount, }) + if i < len(res.ReplayDeleteEvents) { + if deleted := res.ReplayDeleteEvents[i]; deleted != nil && deleted.Pts > 0 && len(deleted.MessageIDs) > 0 { + updates = append(updates, &tg.UpdateDeleteMessages{ + Messages: append([]int(nil), deleted.MessageIDs...), + Pts: deleted.Pts, + PtsCount: deleted.PtsCount, + }) + if deleted.Date > date { + date = deleted.Date + } + } + } if date == 0 { date = event.Date } diff --git a/internal/rpc/messages_forward_rpc_test.go b/internal/rpc/messages_forward_rpc_test.go index 82be8a42..4ccc7e36 100644 --- a/internal/rpc/messages_forward_rpc_test.go +++ b/internal/rpc/messages_forward_rpc_test.go @@ -241,6 +241,88 @@ func TestMessagesForwardMessagesLoadsPrivateSourcesInSingleBatch(t *testing.T) { } } +func TestMessagesForwardMessagesChannelReplayDoesNotRepeatRealtimePayload(t *testing.T) { + ctx := context.Background() + users := memory.NewUserStore() + owner, err := users.Create(ctx, domain.User{AccessHash: 51, Phone: "15550004011", FirstName: "Owner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + source, err := users.Create(ctx, domain.User{AccessHash: 52, Phone: "15550004012", FirstName: "Source"}) + if err != nil { + t.Fatalf("create source: %v", err) + } + channels := appchannels.NewService(memory.NewChannelStore()) + created, err := channels.CreateChannel(ctx, owner.ID, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, + Title: "Forward Replay", + Megagroup: true, + Date: 1700001210, + }) + if err != nil { + t.Fatalf("create target channel: %v", err) + } + messages := &captureMessages{ + getMessagesListed: true, + list: domain.MessageList{Messages: []domain.Message{ + {ID: 7, OwnerUserID: owner.ID, Peer: domain.Peer{Type: domain.PeerTypeUser, ID: source.ID}, From: domain.Peer{Type: domain.PeerTypeUser, ID: source.ID}, Date: 1700001207, Body: "seven"}, + {ID: 5, OwnerUserID: owner.ID, Peer: domain.Peer{Type: domain.PeerTypeUser, ID: source.ID}, From: domain.Peer{Type: domain.PeerTypeUser, ID: source.ID}, Date: 1700001205, Body: "five"}, + }}, + } + sessions := &captureSessions{} + r := New(Config{}, Deps{ + Users: appusers.NewService(users), + Messages: messages, + Channels: channels, + Sessions: sessions, + }, zaptest.NewLogger(t), clock.System) + reqCtx := WithSessionID(WithRawAuthKeyID(WithUserID(ctx, owner.ID), [8]byte{1}), 71) + to := &tg.InputPeerChannel{ChannelID: created.Channel.ID, AccessHash: created.Channel.AccessHash} + from := &tg.InputPeerUser{UserID: source.ID, AccessHash: source.AccessHash} + + firstReq := &tg.MessagesForwardMessagesRequest{FromPeer: from, ToPeer: to, ID: []int{7}, RandomID: []int64{7001}} + if _, err := r.onMessagesForwardMessages(reqCtx, firstReq); err != nil { + t.Fatalf("first forward: %v", err) + } + firstPushes := len(sessions.pushedUserIDs()) + if firstPushes == 0 { + t.Fatal("first forward produced no realtime payload") + } + + if _, err := r.onMessagesForwardMessages(reqCtx, firstReq); err != nil { + t.Fatalf("full replay: %v", err) + } + if got := len(sessions.pushedUserIDs()); got != firstPushes { + t.Fatalf("full replay realtime pushes = %d, want unchanged %d", got, firstPushes) + } + + mixedReq := &tg.MessagesForwardMessagesRequest{ + FromPeer: from, + ToPeer: to, + ID: []int{7, 5}, + RandomID: []int64{7001, 5001}, + } + if _, err := r.onMessagesForwardMessages(reqCtx, mixedReq); err != nil { + t.Fatalf("mixed replay: %v", err) + } + if got := len(sessions.pushedUserIDs()); got != firstPushes+1 { + t.Fatalf("mixed replay realtime pushes = %d, want %d", got, firstPushes+1) + } + updates, ok := sessions.lastUserPush().(*tg.Updates) + if !ok { + t.Fatalf("mixed replay realtime payload = %T, want *tg.Updates", sessions.lastUserPush()) + } + newMessages := 0 + for _, update := range updates.Updates { + if _, ok := update.(*tg.UpdateNewChannelMessage); ok { + newMessages++ + } + } + if newMessages != 1 { + t.Fatalf("mixed replay realtime new-message updates = %d, want only newly committed item", newMessages) + } +} + func TestMessagesForwardMessagesInfersPrivateSourceFromInputPeerEmpty(t *testing.T) { const ( ownerID = int64(1780243210) diff --git a/internal/rpc/messages_monoforum.go b/internal/rpc/messages_monoforum.go index 62276089..5034da1a 100644 --- a/internal/rpc/messages_monoforum.go +++ b/internal/rpc/messages_monoforum.go @@ -191,7 +191,7 @@ func (r *Router) monoforumReplyTargetPeer(userID int64, input tg.InputReplyToCla // sendMonoforumMessage 处理向频道私信(monoforum)发送:订阅者发到自己的子会话,管理员回复到目标订阅者。 // saved_peer 来自 reply_to 的 monoforum_peer_id;管理员可写任意订阅者子会话,普通订阅者只能写自己的。 -func (r *Router) sendMonoforumMessage(ctx context.Context, userID int64, peer domain.Peer, req *tg.MessagesSendMessageRequest) (tg.UpdatesClass, error) { +func (r *Router) sendMonoforumMessage(ctx context.Context, userID int64, peer domain.Peer, req *tg.MessagesSendMessageRequest, fingerprint []byte, preflighted bool) (tg.UpdatesClass, error) { if r.deps.Channels == nil { return nil, notImplementedErr() } @@ -215,13 +215,15 @@ func (r *Router) sendMonoforumMessage(ctx context.Context, userID int64, peer do return nil, replyToMonoforumPeerInvalidErr() } res, err := r.deps.Channels.SendMonoforumMessage(ctx, domain.SendMonoforumMessageRequest{ - MonoforumID: mono.ID, - SenderUserID: userID, - SavedPeer: savedPeer, - RandomID: req.RandomID, - Message: req.Message, - Entities: domainMessageEntities(req.Entities), - Date: int(r.clock.Now().Unix()), + MonoforumID: mono.ID, + SenderUserID: userID, + SavedPeer: savedPeer, + RandomID: req.RandomID, + IdempotencyFingerprint: fingerprint, + IdempotencyPreflighted: preflighted, + Message: req.Message, + Entities: domainMessageEntities(req.Entities), + Date: int(r.clock.Now().Unix()), }) if err != nil { return nil, messageSendErr(err) @@ -232,7 +234,7 @@ func (r *Router) sendMonoforumMessage(ctx context.Context, userID int64, peer do // monoforumSendUpdates 给发送者构造回声 Updates:updateMessageID(关联 random_id)+ updateNewChannelMessage // (monoforum 走 channel pts)。另一方经 monoforum 频道的 getChannelDifference 收取该 durable 事件。 func (r *Router) monoforumSendUpdates(ctx context.Context, userID int64, mono domain.Channel, savedPeer domain.Peer, res domain.SendChannelMessageResult) tg.UpdatesClass { - updates := make([]tg.UpdateClass, 0, 2) + updates := make([]tg.UpdateClass, 0, 3) if res.Message.RandomID != 0 { updates = append(updates, &tg.UpdateMessageID{ID: res.Message.ID, RandomID: res.Message.RandomID}) } @@ -243,10 +245,19 @@ func (r *Router) monoforumSendUpdates(ctx context.Context, userID int64, mono do newMsg.Message = &tg.MessageEmpty{ID: res.Message.ID} } updates = append(updates, newMsg) + date := int(r.clock.Now().Unix()) + if res.Duplicate && res.ReplayDeleteEvent != nil { + if deleted := tgChannelUpdate(userID, *res.ReplayDeleteEvent); deleted != nil { + updates = append(updates, deleted) + } + if res.ReplayDeleteEvent.Date > date { + date = res.ReplayDeleteEvent.Date + } + } return &tg.Updates{ Updates: updates, Chats: r.monoforumChats(ctx, userID, mono), Users: r.monoforumSubscriberUsers(ctx, userID, []domain.MonoforumDialog{{SavedPeer: savedPeer}}, []domain.ChannelMessage{res.Message}), - Date: int(r.clock.Now().Unix()), + Date: date, } } diff --git a/internal/rpc/messages_pin.go b/internal/rpc/messages_pin.go index 6d001498..3aa4fe1c 100644 --- a/internal/rpc/messages_pin.go +++ b/internal/rpc/messages_pin.go @@ -18,7 +18,6 @@ func (r *Router) updatePrivatePinnedMessage(ctx context.Context, userID int64, p if r.deps.Messages == nil { return nil, notImplementedErr() } - authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) res, err := r.deps.Messages.PinPrivateMessage(ctx, userID, domain.PinPrivateMessageRequest{ OwnerUserID: userID, @@ -28,7 +27,7 @@ func (r *Router) updatePrivatePinnedMessage(ctx context.Context, userID int64, p PmOneside: req.PmOneside, Silent: req.Silent, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { @@ -81,7 +80,6 @@ func (r *Router) sendPinServiceMessage(ctx context.Context, userID, peerUserID i if err != nil { return domain.SendPrivateTextResult{}, err } - authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) return r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{ SenderUserID: userID, @@ -96,7 +94,7 @@ func (r *Router) sendPinServiceMessage(ctx context.Context, userID, peerUserID i // 翻译为对端视角；客户端凭此渲染"X 置顶了「…」"预览。 ReplyTo: &domain.MessageReply{MessageID: pinnedBoxID, Peer: domain.Peer{Type: domain.PeerTypeUser, ID: peerUserID}}, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) @@ -121,7 +119,7 @@ func (r *Router) unpinAllPrivateMessages(ctx context.Context, userID int64, peer OwnerUserID: userID, Peer: peer, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { diff --git a/internal/rpc/messages_quick_replies.go b/internal/rpc/messages_quick_replies.go index 34571474..fa42a854 100644 --- a/internal/rpc/messages_quick_replies.go +++ b/internal/rpc/messages_quick_replies.go @@ -198,7 +198,6 @@ func (r *Router) onMessagesSendQuickReplyMessages(ctx context.Context, req *tg.M return nil, err } sessionID, _ := SessionIDFrom(ctx) - authKeyID, _ := AuthKeyIDFrom(ctx) res := domain.ForwardPrivateMessagesResult{OwnerUserID: userID} now := int(r.clock.Now().Unix()) for i, template := range list.Messages { @@ -209,7 +208,7 @@ func (r *Router) onMessagesSendQuickReplyMessages(ctx context.Context, req *tg.M Message: template.Message, Entities: append([]domain.MessageEntity(nil), template.Entities...), Date: now, - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) @@ -361,7 +360,7 @@ func (r *Router) recordQuickReplyMutation(ctx context.Context, userID int64, mut } authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, _, err := r.deps.Updates.RecordQuickReplyMutation(ctx, authKeyID, userID, mutation, sessionID) + event, _, err := r.deps.Updates.RecordQuickReplyMutation(ctx, authKeyID, userID, mutation, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return domain.UpdateEvent{}, internalErr() } diff --git a/internal/rpc/messages_read.go b/internal/rpc/messages_read.go index c623d777..406e156f 100644 --- a/internal/rpc/messages_read.go +++ b/internal/rpc/messages_read.go @@ -28,7 +28,7 @@ func (r *Router) onMessagesReadMessageContents(ctx context.Context, ids []int) ( OwnerUserID: userID, IDs: ids, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: id, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { diff --git a/internal/rpc/messages_register.go b/internal/rpc/messages_register.go index 3c06529e..bf525016 100644 --- a/internal/rpc/messages_register.go +++ b/internal/rpc/messages_register.go @@ -330,7 +330,7 @@ func (r *Router) registerMessages(d *tg.ServerDispatcher) { Peer: peer, MaxID: req.MaxID, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: id, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, }) if err != nil { diff --git a/internal/rpc/messages_saved_dialogs.go b/internal/rpc/messages_saved_dialogs.go index 13d068ff..f1c6b878 100644 --- a/internal/rpc/messages_saved_dialogs.go +++ b/internal/rpc/messages_saved_dialogs.go @@ -177,7 +177,7 @@ func (r *Router) onMessagesToggleSavedDialogPin(ctx context.Context, req *tg.Mes if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordSavedDialogPinned(ctx, authKeyID, userID, peers[0], pinned, sessionID) + event, state, err := r.deps.Updates.RecordSavedDialogPinned(ctx, authKeyID, userID, peers[0], pinned, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } @@ -223,7 +223,7 @@ func (r *Router) onMessagesReorderPinnedSavedDialogs(ctx context.Context, req *t if r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - event, state, err := r.deps.Updates.RecordPinnedSavedDialogs(ctx, authKeyID, userID, peers, sessionID) + event, state, err := r.deps.Updates.RecordPinnedSavedDialogs(ctx, authKeyID, userID, peers, rawAuthKeyIDForOrigin(ctx), sessionID) if err != nil { return false, internalErr() } diff --git a/internal/rpc/messages_send.go b/internal/rpc/messages_send.go index 46b7bbf4..77d7697a 100644 --- a/internal/rpc/messages_send.go +++ b/internal/rpc/messages_send.go @@ -44,11 +44,6 @@ func (r *Router) onMessagesSendMessage(ctx context.Context, req *tg.MessagesSend sendErr = err return nil, sendErr } - // 消息特效：仅接受 catalog 内的合法 effect id（非法 → EFFECT_ID_INVALID，官方行为）。 - if r.messageEffectInvalid(ctx, req.Effect) { - sendErr = effectIDInvalidErr() - return nil, sendErr - } userID, _, err := r.currentUserID(ctx) if err != nil { sendErr = internalErr() @@ -58,25 +53,73 @@ func (r *Router) onMessagesSendMessage(ctx context.Context, req *tg.MessagesSend sendErr = peerIDInvalidErr() return nil, sendErr } - if err := r.checkSendRateLimit(ctx, userID, 1); err != nil { - sendErr = err + peer, ok := r.domainPeerFromInputPeer(userID, req.Peer) + if !ok || peer.ID == 0 { + sendErr = peerIDInvalidErr() return nil, sendErr } - peer, err := r.checkedDomainPeerFromInputPeer(ctx, userID, req.Peer) + idempotencyFingerprint, err := sendMessageIdempotencyFingerprint(req) if err != nil { - sendErr = err + sendErr = internalErr() return nil, sendErr } // 频道私信(monoforum):仅当 reply_to 带 monoforum_peer_id 时走专用发送路径(普通发送恒不带, // 故此 gate 对普通发送零额外成本)。peer 解析为 monoforum 频道时按订阅者子会话发送。 if monoforumReplyPresent(req.ReplyTo) { - updates, err := r.sendMonoforumMessage(ctx, userID, peer, req) + savedPeer, valid := r.monoforumReplyTargetPeer(userID, req.ReplyTo) + if !valid || savedPeer.Type != domain.PeerTypeUser || savedPeer.ID == 0 { + sendErr = replyToMonoforumPeerInvalidErr() + return nil, sendErr + } + replay, err := r.lookupChannelSendReplay(ctx, userID, peer.ID, savedPeer, req.RandomID, idempotencyFingerprint) + if err != nil { + sendErr = err + return nil, err + } + if replay.found { + duplicate = true + return r.monoforumSendUpdates(ctx, userID, replay.channel.Channel, savedPeer, replay.channel), nil + } + if err := r.checkSendRateLimit(ctx, userID, 1); err != nil { + sendErr = err + return nil, sendErr + } + checkedPeer, err := r.checkedDomainPeerFromInputPeer(ctx, userID, req.Peer) + if err != nil { + sendErr = err + return nil, sendErr + } + updates, err := r.sendMonoforumMessage(ctx, userID, checkedPeer, req, idempotencyFingerprint, replay.checked) if err != nil { sendErr = err return nil, sendErr } return updates, nil } + replay, err := r.lookupOutgoingReplay(ctx, userID, peer, req.RandomID, idempotencyFingerprint) + if err != nil { + sendErr = err + return nil, err + } + if replay.found { + duplicate = true + return r.outgoingReplayUpdates(ctx, userID, peer, req.RandomID, replay), nil + } + // Mutable catalog state, rate accounting and access checks are intentionally after exact + // replay lookup: a committed send remains acknowledgeable after those states change. + if r.messageEffectInvalid(ctx, req.Effect) { + sendErr = effectIDInvalidErr() + return nil, sendErr + } + if err := r.checkSendRateLimit(ctx, userID, 1); err != nil { + sendErr = err + return nil, sendErr + } + peer, err = r.checkedDomainPeerFromInputPeer(ctx, userID, req.Peer) + if err != nil { + sendErr = err + return nil, sendErr + } // reply_markup（bot inline keyboard）：仅 bot 账号发送被接受+校验；非 bot 静默丢弃。 // 仅在请求携带 markup 时才查 is_bot，避免普通发送多打一次查询。 var replyMarkup *domain.MessageReplyMarkup @@ -114,16 +157,18 @@ func (r *Router) onMessagesSendMessage(ctx context.Context, req *tg.MessagesSend } if req.ScheduleDate != 0 && !scheduleDateIsImmediate(req.ScheduleDate, int(r.clock.Now().Unix())) { updates, err := r.scheduleOutgoing(ctx, userID, peer, outgoingSend{ - randomID: req.RandomID, - message: req.Message, - entities: req.Entities, - media: previewMedia, - silent: req.Silent, - noforwards: req.Noforwards, - replyToInput: req.ReplyTo, - sendAsInput: req.SendAs, - clearDraft: req.ClearDraft, - richMessage: richMessage, + randomID: req.RandomID, + idempotencyFingerprint: idempotencyFingerprint, + idempotencyPreflighted: replay.checked, + message: req.Message, + entities: req.Entities, + media: previewMedia, + silent: req.Silent, + noforwards: req.Noforwards, + replyToInput: req.ReplyTo, + sendAsInput: req.SendAs, + clearDraft: req.ClearDraft, + richMessage: richMessage, }, req.ScheduleDate, req.ScheduleRepeatPeriod) if err != nil { sendErr = err @@ -132,18 +177,20 @@ func (r *Router) onMessagesSendMessage(ctx context.Context, req *tg.MessagesSend return updates, nil } updates, dup, err := r.sendOutgoing(ctx, userID, peer, outgoingSend{ - randomID: req.RandomID, - message: req.Message, - entities: req.Entities, - media: previewMedia, - silent: req.Silent, - noforwards: req.Noforwards, - replyToInput: req.ReplyTo, - sendAsInput: req.SendAs, - clearDraft: req.ClearDraft, - replyMarkup: replyMarkup, - richMessage: richMessage, - effect: req.Effect, + randomID: req.RandomID, + idempotencyFingerprint: idempotencyFingerprint, + idempotencyPreflighted: replay.checked, + message: req.Message, + entities: req.Entities, + media: previewMedia, + silent: req.Silent, + noforwards: req.Noforwards, + replyToInput: req.ReplyTo, + sendAsInput: req.SendAs, + clearDraft: req.ClearDraft, + replyMarkup: replyMarkup, + richMessage: richMessage, + effect: req.Effect, }) duplicate = dup if err != nil { @@ -159,6 +206,8 @@ func messageSendErr(err error) error { return frozenMethodInvalidErr() case errors.Is(err, domain.ErrReplyMessageIDInvalid): return replyMessageIDInvalidErr() + case errors.Is(err, domain.ErrMessageRandomIDDuplicate): + return randomIDDuplicateErr() default: return internalErr() } @@ -409,6 +458,33 @@ func tgPrivateMessageUpdates(event domain.UpdateEvent, msg domain.Message, rando } } +// tgPrivateSendResultUpdates returns a complete send acknowledgement for exact +// random_id replays. DrKLO requires UpdateNewMessage in an Updates response to +// transition its local pending message to SENT. Visible edited messages use the +// current snapshot; deleted messages use the immutable first snapshot followed +// by the already-durable delete event, so the acknowledgement cannot become a +// permanent resurrection. No replay allocates pts or emits fan-out. +func tgPrivateSendResultUpdates(res domain.SendPrivateTextResult, randomID int64, includeMessageIDForNew bool, users []tg.UserClass, chats []tg.ChatClass) *tg.Updates { + if !res.Duplicate { + return tgPrivateMessageUpdates(res.SenderEvent, res.SenderMessage, randomID, includeMessageIDForNew, users, chats) + } + if randomID == 0 { + randomID = res.SenderMessage.RandomID + } + out := tgPrivateMessageUpdates(res.SenderEvent, res.SenderMessage, randomID, randomID != 0, users, chats) + if event := res.ReplayDeleteEvent; event != nil && event.Pts > 0 && len(event.MessageIDs) > 0 { + out.Updates = append(out.Updates, &tg.UpdateDeleteMessages{ + Messages: append([]int(nil), event.MessageIDs...), + Pts: event.Pts, + PtsCount: event.PtsCount, + }) + if event.Date > out.Date { + out.Date = event.Date + } + } + return out +} + func (r *Router) usersForMessageUpdate(ctx context.Context, ownerUserID int64, msg domain.Message) []tg.UserClass { seen := make(map[int64]struct{}, 2) users := make([]tg.UserClass, 0, 2) diff --git a/internal/rpc/messages_send_rpc_test.go b/internal/rpc/messages_send_rpc_test.go index 6804f38c..f2ff5c46 100644 --- a/internal/rpc/messages_send_rpc_test.go +++ b/internal/rpc/messages_send_rpc_test.go @@ -52,6 +52,9 @@ func TestMessagesSendMessageReturnsUpdateAndRecordsOwnerContext(t *testing.T) { if messages.sendUserID != sender.ID || messages.sendReq.SenderUserID != sender.ID || messages.sendReq.RecipientUserID != recipient.ID || messages.sendReq.OriginSessionID != 77 { t.Fatalf("send context = user %d req %+v, want sender/recipient/session", messages.sendUserID, messages.sendReq) } + if len(messages.sendReq.IdempotencyFingerprint) != 32 { + t.Fatalf("idempotency fingerprint length = %d, want SHA-256", len(messages.sendReq.IdempotencyFingerprint)) + } if len(messages.sendReq.Entities) != 2 || messages.sendReq.Entities[0].Type != domain.MessageEntityBold { t.Fatalf("entities = %+v, want bold and formatted date converted to domain", messages.sendReq.Entities) } diff --git a/internal/rpc/messages_todos.go b/internal/rpc/messages_todos.go index c18e8c8e..fc89b460 100644 --- a/internal/rpc/messages_todos.go +++ b/internal/rpc/messages_todos.go @@ -243,7 +243,6 @@ func (r *Router) mutateTodoMedia(ctx context.Context, inputPeer tg.InputPeerClas return nil, peerIDInvalidErr() } sessionID, _ := SessionIDFrom(ctx) - authKeyID, _ := AuthKeyIDFrom(ctx) res, err := r.deps.Messages.EditMessage(ctx, userID, domain.EditMessageRequest{ OwnerUserID: userID, Peer: peer, @@ -252,7 +251,7 @@ func (r *Router) mutateTodoMedia(ctx context.Context, inputPeer tg.InputPeerClas Entities: current.entities, Media: newMedia, EditDate: now, - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, AllowTodoParticipantMutation: participantEdit, }) diff --git a/internal/rpc/outbox_dispatcher.go b/internal/rpc/outbox_dispatcher.go index 1eed45f7..9d9dd57c 100644 --- a/internal/rpc/outbox_dispatcher.go +++ b/internal/rpc/outbox_dispatcher.go @@ -18,7 +18,10 @@ import ( const ( defaultOutboxBatch = 100 defaultOutboxInterval = 200 * time.Millisecond - defaultOutboxWorkers = 1 + defaultOutboxWorkers = 4 + // outboxLogicalShards 是稳定 user→lane 哈希空间。它不随运行时 worker 数变化， + // worker 只独占其中一组 shard，保证同一用户始终单 lane 串行、不同用户可并行。 + outboxLogicalShards = store.DispatchOutboxLogicalShards // defaultOutboxMaxIdleInterval 是空闲退避上界：连续空 claim 时把轮询间隔从 interval // 指数退避到此上界，削减「无消息时也每 200ms 查一次 DB」的空转；一旦有活就立刻复位到 // interval。代价是「长时间静默后的第一条消息」最多多等这个上界——退避只在持续空闲时触及， @@ -138,25 +141,61 @@ func (d *OutboxDispatcher) Run(ctx context.Context) { if d == nil || d.events == nil || d.outbox == nil || d.sessions == nil { return } - workers := d.workers - if workers < 1 { - workers = 1 - } + claimer, sharded := d.outbox.(shardedOutboxClaimer) + workers := normalizedOutboxWorkers(d.workers, sharded) var wg sync.WaitGroup wg.Add(workers) for i := 0; i < workers; i++ { + var shardIDs []int + if sharded { + shardIDs = logicalShardsForWorker(i, workers) + } go func() { defer wg.Done() - d.runWorker(ctx) + d.runWorker(ctx, claimer, shardIDs) }() } wg.Wait() } +func normalizedOutboxWorkers(workers int, sharded bool) int { + if workers < 1 { + workers = 1 + } + if !sharded { + // 测试替身或旧 store 没有 user shard claim 时，强制单 worker，避免同一用户 + // 的不同 pts 被并行领取。 + return 1 + } + // 多于 logical shard 的 worker 没有独占 lane。若让空 shard worker 回退全局 + // claim，会与全部分片 worker 重叠，因此必须在启动前钳制。 + if workers > outboxLogicalShards { + return outboxLogicalShards + } + return workers +} + +func logicalShardsForWorker(worker, workers int) []int { + if workers <= 0 || worker < 0 || worker >= workers { + return nil + } + out := make([]int, 0, (outboxLogicalShards+workers-1)/workers) + for shard := worker; shard < outboxLogicalShards; shard += workers { + out = append(out, shard) + } + return out +} + // runWorker 是单个 claim 循环；多 worker 靠 ClaimPending 的 SKIP LOCKED 互不重叠。 // 空闲指数退避（interval→maxIdleInterval），claim 到事件即复位到 interval：有活快、无活省。 -func (d *OutboxDispatcher) runWorker(ctx context.Context) { - runIdleBackoffLoop(ctx, d.interval, d.maxIdleInterval, d.DispatchOnce) +func (d *OutboxDispatcher) runWorker(ctx context.Context, claimer shardedOutboxClaimer, shardIDs []int) { + dispatch := d.DispatchOnce + if claimer != nil && len(shardIDs) > 0 { + dispatch = func(ctx context.Context) bool { + return d.dispatchOnceShards(ctx, claimer, shardIDs) + } + } + runIdleBackoffLoop(ctx, d.interval, d.maxIdleInterval, dispatch) } // batchEventLoader 是 UpdateEventStore 的可选批量能力：一次取多条 (user,pts) 事件。 @@ -169,11 +208,24 @@ type batchOutboxMarker interface { MarkDeliveredBatch(ctx context.Context, items []store.DispatchOutboxItem) error } +type shardedOutboxClaimer interface { + ClaimPendingShards(ctx context.Context, shardCount int, shardIDs []int, limit int) ([]store.DispatchOutboxItem, error) +} + // DispatchOnce claim 一批 outbox 并投递，测试可直接调用。返回本次是否 claim 到事件， // 供 runWorker 决定快轮询还是空闲退避。 // store 同时具备批量取事件 + 批量标记能力时走批量路径（每批 ~3 次 PG 往返），否则逐条回退。 func (d *OutboxDispatcher) DispatchOnce(ctx context.Context) bool { items, err := d.outbox.ClaimPending(ctx, d.batch) + return d.dispatchClaimed(ctx, items, err) +} + +func (d *OutboxDispatcher) dispatchOnceShards(ctx context.Context, claimer shardedOutboxClaimer, shardIDs []int) bool { + items, err := claimer.ClaimPendingShards(ctx, outboxLogicalShards, shardIDs, d.batch) + return d.dispatchClaimed(ctx, items, err) +} + +func (d *OutboxDispatcher) dispatchClaimed(ctx context.Context, items []store.DispatchOutboxItem, err error) bool { if err != nil { d.log.Warn("claim dispatch outbox", zap.Error(err)) return false @@ -189,8 +241,14 @@ func (d *OutboxDispatcher) DispatchOnce(ctx context.Context) bool { return true } } + blockedUsers := make(map[int64]struct{}) for _, item := range items { - d.dispatchItem(ctx, item) + if _, blocked := blockedUsers[item.TargetUserID]; blocked { + continue + } + if !d.dispatchItem(ctx, item) { + blockedUsers[item.TargetUserID] = struct{}{} + } } return true } @@ -223,8 +281,14 @@ func (d *OutboxDispatcher) dispatchBatch(ctx context.Context, items []store.Disp if err != nil { // 批量取失败则整批回退逐条路径，让每条各自重试/标失败，不丢进度。 d.log.Warn("batch load dispatch events", zap.Error(err)) + blockedUsers := make(map[int64]struct{}) for _, item := range items { - d.dispatchItem(ctx, item) + if _, blocked := blockedUsers[item.TargetUserID]; blocked { + continue + } + if !d.dispatchItem(ctx, item) { + blockedUsers[item.TargetUserID] = struct{}{} + } } return } @@ -235,10 +299,15 @@ func (d *OutboxDispatcher) dispatchBatch(ctx context.Context, items []store.Disp start := time.Now() ready := make([]outboxDispatchReady, 0, len(items)) requests := make([]OutboxUpdateRequest, 0, len(items)) + blockedUsers := make(map[int64]struct{}) for _, item := range items { + if _, blocked := blockedUsers[item.TargetUserID]; blocked { + continue + } event, ok := byKey[outboxEventKey{item.TargetUserID, item.Pts}] if !ok { d.markDispatchFailed(ctx, item, errMissingOutboxEvent) + blockedUsers[item.TargetUserID] = struct{}{} continue } ready = append(ready, outboxDispatchReady{item: item}) @@ -246,14 +315,19 @@ func (d *OutboxDispatcher) dispatchBatch(ctx context.Context, items []store.Disp } builtUpdates := d.buildOutboxUpdates(ctx, requests) delivered := make([]store.DispatchOutboxItem, 0, len(items)) + clear(blockedUsers) for i, entry := range ready { item := entry.item + if _, blocked := blockedUsers[item.TargetUserID]; blocked { + continue + } update := builtUpdates[i] if update == nil { delivered = append(delivered, item) continue } if _, retriable, err := d.pushOutboxUpdate(ctx, item, update); err != nil { + blockedUsers[item.TargetUserID] = struct{}{} if retriable { // 出站队列拥塞：留 dispatching 行靠租约过期重投，不计入 attempts 升级。 // 不加入 delivered，故不会被 MarkDeliveredBatch 删除。 @@ -271,7 +345,7 @@ func (d *OutboxDispatcher) dispatchBatch(ctx context.Context, items []store.Disp // 批量标记失败则逐条标记，避免整批已投递却卡在 dispatching 等租约过期重投。 d.log.Warn("mark dispatch delivered batch", zap.Error(err)) for _, item := range delivered { - if markErr := d.outbox.MarkDelivered(ctx, item.TargetUserID, item.ID); markErr != nil { + if markErr := d.outbox.MarkDelivered(ctx, item); markErr != nil { d.log.Warn("mark dispatch delivered", zap.Int64("target_user_id", item.TargetUserID), zap.Int64("outbox_id", item.ID), zap.Error(markErr)) } } @@ -286,25 +360,25 @@ type outboxDispatchReady struct { item store.DispatchOutboxItem } -func (d *OutboxDispatcher) dispatchItem(ctx context.Context, item store.DispatchOutboxItem) { +func (d *OutboxDispatcher) dispatchItem(ctx context.Context, item store.DispatchOutboxItem) bool { start := time.Now() events, err := d.events.ListAfter(ctx, item.TargetUserID, item.Pts-1, 1) if err != nil { d.markDispatchFailed(ctx, item, err) - return + return false } if len(events) == 0 || events[0].Pts != item.Pts { d.markDispatchFailed(ctx, item, errMissingOutboxEvent) - return + return false } update := d.buildOutboxUpdate(ctx, item, events[0]) if update == nil { - if err := d.outbox.MarkDelivered(ctx, item.TargetUserID, item.ID); err != nil { + if err := d.outbox.MarkDelivered(ctx, item); err != nil { d.log.Warn("mark noop dispatch delivered", zap.Int64("target_user_id", item.TargetUserID), zap.Int64("outbox_id", item.ID), zap.Error(err)) - return + return false } d.metrics.OutboxDelivered(time.Since(start)) - return + return true } sent, retriable, err := d.pushOutboxUpdate(ctx, item, update) if err != nil { @@ -316,14 +390,14 @@ func (d *OutboxDispatcher) dispatchItem(ctx context.Context, item store.Dispatch zap.Int64("outbox_id", item.ID), zap.Int("pts", item.Pts), ) - return + return false } d.markDispatchFailed(ctx, item, err) - return + return false } - if err := d.outbox.MarkDelivered(ctx, item.TargetUserID, item.ID); err != nil { + if err := d.outbox.MarkDelivered(ctx, item); err != nil { d.log.Warn("mark dispatch delivered", zap.Int64("target_user_id", item.TargetUserID), zap.Int64("outbox_id", item.ID), zap.Error(err)) - return + return false } d.metrics.OutboxDelivered(time.Since(start)) d.log.Debug("dispatch outbox delivered", @@ -332,6 +406,7 @@ func (d *OutboxDispatcher) dispatchItem(ctx context.Context, item store.Dispatch zap.Int("pts", item.Pts), zap.Int("sessions", sent), ) + return true } func (d *OutboxDispatcher) buildOutboxUpdate(ctx context.Context, item store.DispatchOutboxItem, event domain.UpdateEvent) *tg.Updates { @@ -364,19 +439,19 @@ func (d *OutboxDispatcher) buildOutboxUpdates(ctx context.Context, requests []Ou } // pushOutboxUpdate 投递一条 outbox update，返回 (送达的在线 session 数, 是否可重试, err)。 -// best-effort 路径（pushTimeout>0）的失败只可能是出站队列拥塞（慢消费者入队超时），属暂时性、 -// 可重试：调用方应保留 dispatching 行靠租约过期重投，而非计入 attempts 升级为 failed。 -// 可靠路径的失败是真实投递错误，retriable=false，按原逻辑退避升级。 +// 生产 SessionManager 会把 queue-full/closed 慢连接摘除并按离线处理，因此 best-effort +// 接口剩余的非 context 错误通常是确定性的编码/构造错误，必须进入 failed，不能永久占着 +// dispatching head 靠租约空转。只有 dispatcher shutdown/deadline 属于可重试中断。 func (d *OutboxDispatcher) pushOutboxUpdate(ctx context.Context, item store.DispatchOutboxItem, update *tg.Updates) (sent int, retriable bool, err error) { var zeroAuthKeyID [8]byte if d.pushTimeout > 0 { if scoped, ok := d.sessions.(ScopedBestEffortSessionBinder); ok && item.ExcludeAuthKeyID != zeroAuthKeyID { sent, err = scoped.PushToUserExceptAuthKeySessionBestEffort(ctx, item.TargetUserID, item.ExcludeAuthKeyID, item.ExcludeSessionID, proto.MessageFromServer, update, d.pushTimeout) - return sent, err != nil, err + return sent, outboxPushInterrupted(err), err } if bestEffort, ok := d.sessions.(BestEffortSessionBinder); ok { sent, err = bestEffort.PushToUserExceptSessionBestEffort(ctx, item.TargetUserID, item.ExcludeSessionID, proto.MessageFromServer, update, d.pushTimeout) - return sent, err != nil, err + return sent, outboxPushInterrupted(err), err } } if scoped, ok := d.sessions.(ScopedSessionBinder); ok && item.ExcludeAuthKeyID != zeroAuthKeyID { @@ -387,11 +462,15 @@ func (d *OutboxDispatcher) pushOutboxUpdate(ctx context.Context, item store.Disp return sent, false, err } +func outboxPushInterrupted(err error) bool { + return errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) +} + func (d *OutboxDispatcher) markDispatchFailed(ctx context.Context, item store.DispatchOutboxItem, err error) { if err == nil { err = errMissingOutboxEvent } - if markErr := d.outbox.MarkFailed(ctx, item.TargetUserID, item.ID, err.Error()); markErr != nil { + if markErr := d.outbox.MarkFailed(ctx, item, err.Error()); markErr != nil { d.log.Warn("mark dispatch failed", zap.Int64("target_user_id", item.TargetUserID), zap.Int64("outbox_id", item.ID), diff --git a/internal/rpc/outbox_dispatcher_test.go b/internal/rpc/outbox_dispatcher_test.go index 8e519788..22f244dc 100644 --- a/internal/rpc/outbox_dispatcher_test.go +++ b/internal/rpc/outbox_dispatcher_test.go @@ -460,6 +460,101 @@ func TestOutboxDispatcherOrdersClaimedItemsByUserPts(t *testing.T) { } } +func TestOutboxLogicalShardsAreDisjointAndStable(t *testing.T) { + for _, workers := range []int{1, 2, 4, 7, 64, outboxLogicalShards} { + seen := make([]int, outboxLogicalShards) + for worker := 0; worker < workers; worker++ { + for _, shard := range logicalShardsForWorker(worker, workers) { + if shard < 0 || shard >= outboxLogicalShards { + t.Fatalf("workers=%d worker=%d returned invalid shard %d", workers, worker, shard) + } + seen[shard]++ + } + } + for shard, owners := range seen { + if owners != 1 { + t.Fatalf("workers=%d shard=%d owners=%d, want exactly one", workers, shard, owners) + } + } + } + if got := normalizedOutboxWorkers(8, false); got != 1 { + t.Fatalf("non-sharded workers = %d, want 1", got) + } + if got := normalizedOutboxWorkers(outboxLogicalShards+100, true); got != outboxLogicalShards { + t.Fatalf("overprovisioned workers = %d, want clamp %d", got, outboxLogicalShards) + } +} + +func TestOutboxDispatcherBatchFailureBlocksHigherUserPts(t *testing.T) { + const ( + blockedUser = int64(1000000002) + otherUser = int64(1000000003) + ) + items := []store.DispatchOutboxItem{ + {ID: 12, TargetUserID: blockedUser, Pts: 12, EventType: domain.UpdateEventReadHistoryInbox}, + {ID: 5, TargetUserID: otherUser, Pts: 5, EventType: domain.UpdateEventReadHistoryInbox}, + {ID: 11, TargetUserID: blockedUser, Pts: 11, EventType: domain.UpdateEventReadHistoryInbox}, + } + events := make([]domain.UpdateEvent, 0, len(items)) + for _, item := range items { + events = append(events, outboxReadEvent(item.TargetUserID, item.Pts)) + } + eventStore := &batchEventStore{captureUpdateEventStore: &captureUpdateEventStore{events: events}} + outbox := &batchDispatchOutbox{captureDispatchOutbox: &captureDispatchOutbox{items: items}} + sessions := &selectiveFailOutboxSessions{failUserID: blockedUser, failPts: 11} + dispatcher := NewOutboxDispatcher(eventStore, outbox, sessions, zaptest.NewLogger(t)) + + dispatcher.DispatchOnce(context.Background()) + + wantAttempts := []outboxPushAttempt{{userID: blockedUser, pts: 11}, {userID: otherUser, pts: 5}} + if got := sessions.pushAttempts(); !reflect.DeepEqual(got, wantAttempts) { + t.Fatalf("push attempts = %+v, want %+v (blocked user's pts=12 must not overtake failed pts=11)", got, wantAttempts) + } + if !outbox.failed || len(outbox.deliveredBatch) != 1 || outbox.deliveredBatch[0].TargetUserID != otherUser { + t.Fatalf("outbox failed=%v delivered=%+v, want failed head and only other user delivered", outbox.failed, outbox.deliveredBatch) + } +} + +func TestOutboxDispatcherBatchLoadFallbackStillBlocksHigherUserPts(t *testing.T) { + const ( + blockedUser = int64(1000000004) + otherUser = int64(1000000005) + ) + items := []store.DispatchOutboxItem{ + {ID: 22, TargetUserID: blockedUser, Pts: 22, EventType: domain.UpdateEventReadHistoryInbox}, + {ID: 21, TargetUserID: blockedUser, Pts: 21, EventType: domain.UpdateEventReadHistoryInbox}, + {ID: 6, TargetUserID: otherUser, Pts: 6, EventType: domain.UpdateEventReadHistoryInbox}, + } + events := make([]domain.UpdateEvent, 0, len(items)) + for _, item := range items { + events = append(events, outboxReadEvent(item.TargetUserID, item.Pts)) + } + eventStore := &failingBatchEventStore{captureUpdateEventStore: &captureUpdateEventStore{events: events}} + outbox := &batchDispatchOutbox{captureDispatchOutbox: &captureDispatchOutbox{items: items}} + sessions := &selectiveFailOutboxSessions{failUserID: blockedUser, failPts: 21} + dispatcher := NewOutboxDispatcher(eventStore, outbox, sessions, zaptest.NewLogger(t)) + + dispatcher.DispatchOnce(context.Background()) + + wantAttempts := []outboxPushAttempt{{userID: blockedUser, pts: 21}, {userID: otherUser, pts: 6}} + if got := sessions.pushAttempts(); !reflect.DeepEqual(got, wantAttempts) { + t.Fatalf("fallback push attempts = %+v, want %+v", got, wantAttempts) + } +} + +func outboxReadEvent(userID int64, pts int) domain.UpdateEvent { + return domain.UpdateEvent{ + UserID: userID, + Type: domain.UpdateEventReadHistoryInbox, + Pts: pts, + PtsCount: 1, + Date: 1700000000 + pts, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: 999}, + MaxID: pts, + StillUnreadCount: 0, + } +} + type outboxUsersCall struct { viewerUserID int64 ids []int64 @@ -602,6 +697,34 @@ type orderedOutboxCaptureSessions struct { pushed []int } +type outboxPushAttempt struct { + userID int64 + pts int +} + +type selectiveFailOutboxSessions struct { + captureSessions + failUserID int64 + failPts int + attempts []outboxPushAttempt +} + +func (s *selectiveFailOutboxSessions) PushToUserExceptSession(_ context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) { + pts := 0 + if updates, ok := msg.(*tg.Updates); ok { + pts = firstOutboxUpdatePts(updates) + } + s.attempts = append(s.attempts, outboxPushAttempt{userID: userID, pts: pts}) + if userID == s.failUserID && pts == s.failPts { + return 0, errors.New("injected outbox push failure") + } + return s.captureSessions.PushToUserExceptSession(context.Background(), userID, excludeSessionID, t, msg) +} + +func (s *selectiveFailOutboxSessions) pushAttempts() []outboxPushAttempt { + return append([]outboxPushAttempt(nil), s.attempts...) +} + func (s *orderedOutboxCaptureSessions) PushToUserExceptSession(_ context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) { if updates, ok := msg.(*tg.Updates); ok { s.pushed = append(s.pushed, firstOutboxUpdatePts(updates)) @@ -635,6 +758,33 @@ type batchEventStore struct { batchCursors []store.EventCursor } +type failingBatchEventStore struct { + *captureUpdateEventStore +} + +func (s *failingBatchEventStore) BatchByCursor(context.Context, []store.EventCursor) ([]domain.UpdateEvent, error) { + return nil, errors.New("injected batch event load failure") +} + +func (s *failingBatchEventStore) ListAfter(_ context.Context, userID int64, pts, limit int) ([]domain.UpdateEvent, error) { + if limit <= 0 { + return nil, nil + } + var next domain.UpdateEvent + for _, event := range s.events { + if event.UserID != userID || event.Pts <= pts { + continue + } + if next.Pts == 0 || event.Pts < next.Pts { + next = event + } + } + if next.Pts == 0 { + return nil, nil + } + return []domain.UpdateEvent{next}, nil +} + func (s *batchEventStore) BatchByCursor(_ context.Context, cursors []store.EventCursor) ([]domain.UpdateEvent, error) { s.batchCursors = cursors out := make([]domain.UpdateEvent, 0, len(cursors)) @@ -737,6 +887,8 @@ type captureScopedSessions struct { scopedMu sync.Mutex scopedAuthKeyID [8]byte immediatePush bool + immediateType proto.MessageType + immediateMsg bin.Encoder } func (s *captureScopedSessions) setScopedAuthKeyID(rawAuthKeyID [8]byte) { @@ -757,6 +909,12 @@ func (s *captureScopedSessions) immediatePushSeen() bool { return s.immediatePush } +func (s *captureScopedSessions) immediatePushSnapshot() (proto.MessageType, bin.Encoder) { + s.scopedMu.Lock() + defer s.scopedMu.Unlock() + return s.immediateType, s.immediateMsg +} + func (s *captureScopedSessions) BindAuthKeyForSession(rawAuthKeyID [8]byte, sessionID int64, authKeyID [8]byte) { s.BindAuthKey(sessionID, authKeyID) s.setScopedAuthKeyID(rawAuthKeyID) @@ -790,6 +948,8 @@ func (s *captureScopedSessions) PushToSessionForAuthKeyImmediate(_ context.Conte s.scopedMu.Lock() s.immediatePush = true s.scopedAuthKeyID = rawAuthKeyID + s.immediateType = t + s.immediateMsg = msg s.scopedMu.Unlock() return s.PushToSession(context.Background(), sessionID, t, msg) } @@ -805,14 +965,14 @@ func (s *captureDispatchOutbox) ClaimPending(context.Context, int) ([]store.Disp return items, nil } -func (s *captureDispatchOutbox) MarkDelivered(_ context.Context, targetUserID, id int64) error { +func (s *captureDispatchOutbox) MarkDelivered(_ context.Context, item store.DispatchOutboxItem) error { s.delivered = true - s.deliveredUserID = targetUserID - s.deliveredID = id + s.deliveredUserID = item.TargetUserID + s.deliveredID = item.ID return nil } -func (s *captureDispatchOutbox) MarkFailed(_ context.Context, _ int64, _ int64, lastError string) error { +func (s *captureDispatchOutbox) MarkFailed(_ context.Context, _ store.DispatchOutboxItem, lastError string) error { s.failed = true s.failedError = lastError return nil @@ -869,21 +1029,19 @@ func (m *captureOutboxMetrics) OutboxFailed(error) { m.failed++ } -// queueFullBestEffortSessions 模拟出站队列拥塞：best-effort 推送总是失败（入队超时 / 队列满）。 -type queueFullBestEffortSessions struct { +// interruptedBestEffortSessions 模拟 dispatcher context 到期：该中断可安全靠 lease 重试。 +type interruptedBestEffortSessions struct { *captureSessions attempts int } -func (s *queueFullBestEffortSessions) PushToUserExceptSessionBestEffort(_ context.Context, _ int64, _ int64, _ proto.MessageType, _ bin.Encoder, _ time.Duration) (int, error) { +func (s *interruptedBestEffortSessions) PushToUserExceptSessionBestEffort(_ context.Context, _ int64, _ int64, _ proto.MessageType, _ bin.Encoder, _ time.Duration) (int, error) { s.attempts++ - return 0, errors.New("mtproto outbound queue full") + return 0, context.DeadlineExceeded } -// TestOutboxDispatcherDefersOnPushQueueFull 验证 best-effort 推送因出站队列拥塞失败时，dispatcher -// 既不标记 delivered（任务保留，靠 dispatching 租约过期重投，满足至少一次投递语义），也不标记 -// failed（拥塞不计入 attempts 升级，避免正常满 fan-out 负载把可靠 update 误打成 failed）。 -func TestOutboxDispatcherDefersOnPushQueueFull(t *testing.T) { +// TestOutboxDispatcherDefersOnPushInterruption 验证 shutdown/deadline 不把 lane head 误打 failed。 +func TestOutboxDispatcherDefersOnPushInterruption(t *testing.T) { msg := domain.Message{ ID: 10, OwnerUserID: 1000000002, @@ -909,7 +1067,7 @@ func TestOutboxDispatcherDefersOnPushQueueFull(t *testing.T) { EventType: domain.UpdateEventNewMessage, ExcludeSessionID: 99, }}} - sessions := &queueFullBestEffortSessions{captureSessions: &captureSessions{}} + sessions := &interruptedBestEffortSessions{captureSessions: &captureSessions{}} metrics := &captureOutboxMetrics{} dispatcher := NewOutboxDispatcher(events, outbox, sessions, zaptest.NewLogger(t), WithOutboxPushTimeout(50*time.Millisecond), WithOutboxMetrics(metrics)) dispatcher.DispatchOnce(context.Background()) @@ -918,12 +1076,21 @@ func TestOutboxDispatcherDefersOnPushQueueFull(t *testing.T) { t.Fatalf("best-effort push attempts = %d, want 1（应走 best-effort 推送路径）", sessions.attempts) } if outbox.delivered { - t.Fatalf("outbox delivered=true, want 未投递（拥塞应保留 dispatching 行靠租约重投）") + t.Fatalf("outbox delivered=true, want 未投递（中断应保留 dispatching 行靠租约重投）") } if outbox.failed { - t.Fatalf("outbox failed=true, want 未失败（拥塞不计入 attempts 升级）") + t.Fatalf("outbox failed=true, want 未失败（context 中断不计入 attempts 升级）") } if metrics.failed != 0 { - t.Fatalf("metrics.failed=%d, want 0（拥塞不算投递失败）", metrics.failed) + t.Fatalf("metrics.failed=%d, want 0（context 中断不算投递失败）", metrics.failed) + } +} + +func TestOutboxPushInterruptedRejectsDeterministicErrors(t *testing.T) { + if !outboxPushInterrupted(context.Canceled) || !outboxPushInterrupted(context.DeadlineExceeded) { + t.Fatal("context shutdown/deadline must remain retriable") + } + if outboxPushInterrupted(errors.New("encode update: invalid constructor")) { + t.Fatal("deterministic encoding error must fail the lane head instead of lease-retrying forever") } } diff --git a/internal/rpc/payments_star_gifts.go b/internal/rpc/payments_star_gifts.go index a91871e1..279e6b9b 100644 --- a/internal/rpc/payments_star_gifts.go +++ b/internal/rpc/payments_star_gifts.go @@ -340,7 +340,6 @@ func (r *Router) deliverStarGift(ctx context.Context, senderID, recipientID int6 if err != nil { return domain.SendPrivateTextResult{}, err } - authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) sticker := gift.Sticker media := &domain.MessageMedia{ @@ -368,7 +367,7 @@ func (r *Router) deliverStarGift(ctx context.Context, senderID, recipientID int6 Media: media, Silent: false, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) diff --git a/internal/rpc/phone_conference.go b/internal/rpc/phone_conference.go index 21939b4d..8a442040 100644 --- a/internal/rpc/phone_conference.go +++ b/internal/rpc/phone_conference.go @@ -146,10 +146,11 @@ func (r *Router) onPhoneInviteConferenceCallParticipant(ctx context.Context, req return nil, err } now := int(r.clock.Now().Unix()) + randomID := conferenceInviteRandomID(scope.call.ID, target.ID, now) res, err := r.deps.Messages.SendPrivateText(ctx, scope.userID, domain.SendPrivateTextRequest{ SenderUserID: scope.userID, RecipientUserID: target.ID, - RandomID: conferenceInviteRandomID(scope.call.ID, target.ID, now), + RandomID: randomID, Media: &domain.MessageMedia{ Kind: domain.MessageMediaKindService, ServiceAction: &domain.MessageServiceAction{ @@ -164,13 +165,16 @@ func (r *Router) onPhoneInviteConferenceCallParticipant(ctx context.Context, req }, }, Date: now, - OriginAuthKeyID: authKeyIDFromCtx(ctx), + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionIDFromCtx(ctx), RecipientBlocked: recipientBlocked, }) if err != nil { return nil, messageSendErr(err) } + // Message and invite index currently live in separate stores/transactions. Always + // retry the idempotent invite write, including an exact message replay, so a crash + // after SendPrivateText commit cannot leave an unrecoverable message-without-index. invite, err := r.deps.GroupCalls.CreateConferenceInvite(ctx, domain.GroupCallInvite{ CallID: scope.call.ID, InviterUserID: scope.userID, @@ -184,6 +188,15 @@ func (r *Router) onPhoneInviteConferenceCallParticipant(ctx context.Context, req return nil, groupCallErr(err) } _ = invite + if res.Duplicate { + // The first transaction already created both private boxes and their durable + // update events. Replaying UpdateNewMessage here would reconstruct the service + // message from the intentionally minimal immutable receipt and push it to the + // invitee a second time. Only reconcile the caller's random_id and original pts; + // The invite write above is an idempotent saga repair only; no message/update + // side effect is repeated. + return tgPrivateSendResultUpdates(res, randomID, true, nil, nil), nil + } users := r.tgUsersForIDs(ctx, scope.userID, []int64{scope.userID, target.ID}) out := tgPrivateMessageUpdates(res.SenderEvent, res.SenderMessage, 0, false, users, nil) recipientUsers := r.tgUsersForIDs(ctx, target.ID, []int64{scope.userID, target.ID}) diff --git a/internal/rpc/phone_conference_rpc_test.go b/internal/rpc/phone_conference_rpc_test.go index 4da56b89..4d95ca6c 100644 --- a/internal/rpc/phone_conference_rpc_test.go +++ b/internal/rpc/phone_conference_rpc_test.go @@ -799,6 +799,82 @@ func TestConferenceInviteMessageResolvesInputGroupCallInviteMessage(t *testing.T } } +func TestConferenceInviteExactReplayConfirmsImmutableMessageWithoutFanout(t *testing.T) { + f := newConferenceFixture(t) + aliceCtx := f.userCtx(f.alice, 11) + create, err := f.router.onPhoneCreateConferenceCall(aliceCtx, &tg.PhoneCreateConferenceCallRequest{RandomID: 89}) + if err != nil { + t.Fatalf("create conference: %v", err) + } + call := findUpdate[*tg.UpdateGroupCall](t, create).Call.(*tg.GroupCall) + request := &tg.PhoneInviteConferenceCallParticipantRequest{ + Call: &tg.InputGroupCall{ID: call.ID, AccessHash: call.AccessHash}, + UserID: &tg.InputUser{UserID: f.bob.ID, AccessHash: f.bob.AccessHash}, + } + first, err := f.router.onPhoneInviteConferenceCallParticipant(aliceCtx, request) + if err != nil { + t.Fatalf("first conference invite: %v", err) + } + firstUpdate := findUpdate[*tg.UpdateNewMessage](t, first) + firstMessage, ok := firstUpdate.Message.(*tg.MessageService) + if !ok { + t.Fatalf("first conference message = %T, want MessageService", firstUpdate.Message) + } + bobHistory, err := f.messages.GetHistory(f.ctx, f.bob.ID, domain.MessageFilter{ + HasPeer: true, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: f.alice.ID}, + Limit: 10, + }) + if err != nil || len(bobHistory.Messages) != 1 { + t.Fatalf("bob history after first invite len=%d err=%v, want one", len(bobHistory.Messages), err) + } + bobMessageID := bobHistory.Messages[0].ID + _, firstInvite, found, err := f.group.GetByInviteMessage(f.ctx, f.bob.ID, bobMessageID) + if err != nil || !found { + t.Fatalf("first durable invite = %+v found=%v err=%v", firstInvite, found, err) + } + + f.sessions.reset() + replay, err := f.router.onPhoneInviteConferenceCallParticipant(aliceCtx, request) + if err != nil { + t.Fatalf("replay conference invite: %v", err) + } + replayed := findUpdate[*tg.UpdateNewMessage](t, replay) + replayedMessage, ok := replayed.Message.(*tg.MessageService) + if !ok || replayedMessage.ID != firstMessage.ID || replayed.Pts != firstUpdate.Pts || replayed.PtsCount != firstUpdate.PtsCount { + t.Fatalf("replay message = %+v (%T), want original confirmation %d pts %d/%d", replayed.Message, replayed.Message, firstMessage.ID, firstUpdate.Pts, firstUpdate.PtsCount) + } + mapping := findUpdate[*tg.UpdateMessageID](t, replay) + wantRandomID := conferenceInviteRandomID(call.ID, f.bob.ID, int(f.clock.Now().Unix())) + if mapping.ID != firstMessage.ID || mapping.RandomID != wantRandomID { + t.Fatalf("replay mapping = %+v, want id/random_id %d/%d", mapping, firstMessage.ID, wantRandomID) + } + if records := f.sessions.records(); len(records) != 0 { + t.Fatalf("replay must not fan out another invite, got %+v", records) + } + bobHistory, err = f.messages.GetHistory(f.ctx, f.bob.ID, domain.MessageFilter{ + HasPeer: true, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: f.alice.ID}, + Limit: 10, + }) + if err != nil || len(bobHistory.Messages) != 1 { + t.Fatalf("bob history after replay len=%d err=%v, want original one", len(bobHistory.Messages), err) + } + _, replayInvite, found, err := f.group.GetByInviteMessage(f.ctx, f.bob.ID, bobMessageID) + if err != nil || !found || replayInvite != firstInvite { + t.Fatalf("durable invite after replay = %+v found=%v err=%v, want unchanged %+v", replayInvite, found, err, firstInvite) + } + + conflict := *request + conflict.Video = true + if result, err := f.router.onPhoneInviteConferenceCallParticipant(aliceCtx, &conflict); result != nil || !tgerr.Is(err, "RANDOM_ID_DUPLICATE") { + t.Fatalf("conflicting replay = %+v err=%v, want RANDOM_ID_DUPLICATE", result, err) + } + if records := f.sessions.records(); len(records) != 0 { + t.Fatalf("conflicting replay must not fan out, got %+v", records) + } +} + func TestPhoneDiscardMigrateConferenceCarriesSlug(t *testing.T) { f := newConferenceFixture(t) aliceCtx := f.userCtx(f.alice, 11) diff --git a/internal/rpc/photos.go b/internal/rpc/photos.go index 0de92d93..48066848 100644 --- a/internal/rpc/photos.go +++ b/internal/rpc/photos.go @@ -318,7 +318,6 @@ func (r *Router) sendSuggestedProfilePhotoMessage(ctx context.Context, userID, t if err != nil { return domain.SendPrivateTextResult{}, err } - authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) photoCopy := photo res, err := r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{ @@ -333,7 +332,7 @@ func (r *Router) sendSuggestedProfilePhotoMessage(ctx context.Context, userID, t }, }, Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, + OriginAuthKeyID: rawAuthKeyIDForOrigin(ctx), OriginSessionID: sessionID, RecipientBlocked: recipientBlocked, }) diff --git a/internal/rpc/push.go b/internal/rpc/push.go index 4f58a50c..7a875225 100644 --- a/internal/rpc/push.go +++ b/internal/rpc/push.go @@ -14,7 +14,7 @@ func (r *Router) pushUserMessage(ctx context.Context, userID int64, logMessage s } sessionID, _ := SessionIDFrom(ctx) if timeout := r.cfg.OutboundPushTimeout; timeout > 0 { - authKeyID, _ := AuthKeyIDFrom(ctx) + authKeyID := rawAuthKeyIDForOrigin(ctx) if scoped, ok := r.deps.Sessions.(ScopedBestEffortSessionBinder); ok { if sent, err := scoped.PushToUserExceptAuthKeySessionBestEffort(ctx, userID, authKeyID, sessionID, proto.MessageFromServer, msg, timeout); err != nil { r.log.Debug(logMessage, zap.Int64("user_id", userID), zap.Int("sent", sent), zap.Duration("timeout", timeout), zap.Error(err)) @@ -33,7 +33,7 @@ func (r *Router) pushUserMessage(ctx context.Context, userID int64, logMessage s } } if scoped, ok := r.scopedSessions(); ok { - authKeyID, _ := AuthKeyIDFrom(ctx) + authKeyID := rawAuthKeyIDForOrigin(ctx) if sent, err := scoped.PushToUserExceptAuthKeySession(ctx, userID, authKeyID, sessionID, proto.MessageFromServer, msg); err != nil { r.log.Debug(logMessage, zap.Int64("user_id", userID), zap.Int("sent", sent), zap.Error(err)) return sent @@ -58,7 +58,7 @@ func (r *Router) pushUserMessageTransient(ctx context.Context, userID int64, log } if transient, ok := r.deps.Sessions.(TransientSessionBinder); ok { sessionID, _ := SessionIDFrom(ctx) - authKeyID, _ := AuthKeyIDFrom(ctx) + authKeyID := rawAuthKeyIDForOrigin(ctx) sent, err := transient.PushToUserTransientExceptAuthKeySession(ctx, userID, authKeyID, sessionID, proto.MessageFromServer, msg, r.cfg.OutboundPushTimeout) if err != nil { r.log.Debug(logMessage, zap.Int64("user_id", userID), zap.Int("sent", sent), zap.Error(err)) diff --git a/internal/rpc/rate_limit.go b/internal/rpc/rate_limit.go index 863c6f82..a857505b 100644 --- a/internal/rpc/rate_limit.go +++ b/internal/rpc/rate_limit.go @@ -2,14 +2,24 @@ package rpc import ( "context" + "crypto/sha256" + "encoding/hex" "strconv" "time" "go.uber.org/zap" + + "telesrv/internal/domain" ) const sendRateLimitKeyPrefix = "messages:send:" +const ( + authCodePhoneRateLimitKeyPrefix = "auth:code:phone-sha256:" + authCodeAuthKeyRateLimitKeyPrefix = "auth:code:raw-auth-key:" + defaultAuthCodeRateWindow = 10 * time.Minute +) + const ( channelDifferenceRateLimitKeyPrefix = "updates:channeldifference:" peerDialogsRateLimitKeyPrefix = "messages:peerdialogs:" @@ -65,3 +75,55 @@ func (r *Router) checkSendRateLimit(ctx context.Context, userID int64, cost int) r.metrics().MessageRateLimited(retryAfter) return floodWaitErr(retryAfter) } + +// checkAuthCodeRateLimit protects the unauthenticated code-issuance path before +// any account lookup or durable 777000 write. Existing and unknown phone numbers +// therefore consume identical budgets and cannot be distinguished through the +// limiter. Plaintext phone numbers are never used as limiter keys or log fields. +func (r *Router) checkAuthCodeRateLimit(ctx context.Context, phone string) error { + if r.deps.Limiter == nil { + return nil + } + normalizedPhone := domain.NormalizePhone(phone) + if !domain.ValidPhone(normalizedPhone) { + return phoneNumberInvalidErr() + } + window := r.cfg.AuthCodeRateWindow + if window <= 0 { + window = defaultAuthCodeRateWindow + } + // Check the connection/auth-key budget first. If that dimension is already + // blocked, changing phone strings cannot create one phone-digest Redis key + // per attempt and bypass the intended cardinality bound. + if limit := r.cfg.AuthCodeAuthKeyRateLimit; limit > 0 { + if rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx); ok && rawAuthKeyID != ([8]byte{}) { + if err := r.checkAuthCodeRateLimitKey(ctx, authCodeAuthKeyRateLimitKeyPrefix+hex.EncodeToString(rawAuthKeyID[:]), limit, window, "raw_auth_key"); err != nil { + return err + } + } + } + if limit := r.cfg.AuthCodePhoneRateLimit; limit > 0 { + digest := sha256.Sum256([]byte(normalizedPhone)) + if err := r.checkAuthCodeRateLimitKey(ctx, authCodePhoneRateLimitKeyPrefix+hex.EncodeToString(digest[:]), limit, window, "phone_digest"); err != nil { + return err + } + } + return nil +} + +func (r *Router) checkAuthCodeRateLimitKey(ctx context.Context, key string, limit int, window time.Duration, dimension string) error { + allowed, retryAfter, err := r.deps.Limiter.AllowN(ctx, key, 1, limit, window) + if err != nil { + return internalErr() + } + if allowed { + return nil + } + if retryAfter <= 0 { + retryAfter = 1 + } + r.log.Debug("auth code issuance rate limited", + zap.String("dimension", dimension), + zap.Int("retry_after", retryAfter)) + return floodWaitErr(retryAfter) +} diff --git a/internal/rpc/request_preflight.go b/internal/rpc/request_preflight.go new file mode 100644 index 00000000..9cb7f195 --- /dev/null +++ b/internal/rpc/request_preflight.go @@ -0,0 +1,137 @@ +package rpc + +import ( + "encoding/binary" + "fmt" + + "github.com/gotd/td/bin" + "github.com/gotd/td/tg" + + appfiles "telesrv/internal/app/files" + "telesrv/internal/domain" +) + +const tlVectorTypeID = uint32(0x1cb5c415) + +type requestVectorPolicy struct { + vectorOffset int + max int + minElemBytes int + tooLong func() error +} + +// requestVectorPolicies mirrors limits already enforced by typed handlers, but does so before +// gotd's generated decoder materializes attacker-controlled interface slices. users.getUsers is +// the one newly introduced cap: TDesktop's four current call sites all send exactly one user. +var requestVectorPolicies = map[uint32]requestVectorPolicy{ + tg.UsersGetUsersRequestTypeID: {vectorOffset: 4, max: 100, minElemBytes: 4, tooLong: inputRequestTooLongErr}, + tg.UsersGetRequirementsToContactRequestTypeID: {vectorOffset: 4, max: maxRequirementsToContactUsers, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.ContactsImportContactsRequestTypeID: {vectorOffset: 4, max: maxContactImportBatch, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.ContactsDeleteContactsRequestTypeID: {vectorOffset: 4, max: maxContactDeleteBatch, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.ContactsEditCloseFriendsRequestTypeID: {vectorOffset: 4, max: maxCloseFriendsCount, minElemBytes: 8, tooLong: limitInvalidErr}, + tg.ContactsSetBlockedRequestTypeID: {vectorOffset: 8, max: maxContactSetBlocked, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.MessagesGetMessagesRequestTypeID: {vectorOffset: 4, max: maxGetMessagesIDs, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.MessagesGetChatsRequestTypeID: {vectorOffset: 4, max: maxGetMessagesIDs, minElemBytes: 8, tooLong: limitInvalidErr}, + tg.MessagesGetPeerDialogsRequestTypeID: {vectorOffset: 4, max: maxDialogInputPeers, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.MessagesReadMessageContentsRequestTypeID: {vectorOffset: 4, max: maxGetMessagesIDs, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.MessagesGetCustomEmojiDocumentsRequestTypeID: {vectorOffset: 4, max: maxEmojiDocuments, minElemBytes: 8, tooLong: limitInvalidErr}, + tg.MessagesDeleteMessagesRequestTypeID: {vectorOffset: 8, max: domain.MaxDeleteMessageIDs, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.MessagesCreateChatRequestTypeID: {vectorOffset: 8, max: 200, minElemBytes: 4, tooLong: limitInvalidErr}, + tg.ChannelsGetChannelsRequestTypeID: {vectorOffset: 4, max: maxGetMessagesIDs, minElemBytes: 4, tooLong: limitInvalidErr}, +} + +func preflightRPCRequest(id uint32, b *bin.Buffer) error { + if b == nil { + return inputRequestInvalidErr() + } + if policy, ok := requestVectorPolicies[id]; ok { + if err := preflightFixedVector(b.Buf, policy); err != nil { + return err + } + } + switch id { + case tg.UploadSaveFilePartRequestTypeID: + return preflightUploadPart(b.Buf, 16, false) + case tg.UploadSaveBigFilePartRequestTypeID: + return preflightUploadPart(b.Buf, 20, true) + default: + return nil + } +} + +func preflightFixedVector(raw []byte, policy requestVectorPolicy) error { + if policy.vectorOffset < 4 || policy.minElemBytes <= 0 || len(raw) < policy.vectorOffset+8 { + return inputRequestInvalidErr() + } + if binary.LittleEndian.Uint32(raw[policy.vectorOffset:]) != tlVectorTypeID { + return inputRequestInvalidErr() + } + count := int64(int32(binary.LittleEndian.Uint32(raw[policy.vectorOffset+4:]))) + if count < 0 { + return inputRequestInvalidErr() + } + remaining := int64(len(raw) - policy.vectorOffset - 8) + // Check the cheapest possible encoding before the policy cap. A forged MaxInt32 count with + // a truncated body is malformed, not merely a large valid request, and is rejected O(1). + if count > remaining/int64(policy.minElemBytes) { + return inputRequestInvalidErr() + } + if count > int64(policy.max) { + if policy.tooLong != nil { + return policy.tooLong() + } + return inputRequestTooLongErr() + } + return nil +} + +func preflightUploadPart(raw []byte, bytesOffset int, big bool) error { + if big { + if len(raw) < 20 { + return inputRequestInvalidErr() + } + totalParts := int32(binary.LittleEndian.Uint32(raw[16:20])) + if totalParts <= 0 || totalParts > appfiles.MaxUploadParts { + return filePartInvalidErr() + } + } + n, encoded, err := tlBytesSizeAt(raw, bytesOffset) + if err != nil { + return inputRequestInvalidErr() + } + if encoded != len(raw)-bytesOffset { + return inputRequestInvalidErr() + } + if n > appfiles.MaxUploadPartBytes { + return filePartTooBigErr() + } + return nil +} + +// tlBytesSizeAt parses a TL bytes prefix without copying the payload. encoded includes prefix, +// payload and 4-byte padding. +func tlBytesSizeAt(raw []byte, offset int) (n, encoded int, err error) { + if offset < 0 || offset >= len(raw) { + return 0, 0, fmt.Errorf("bytes prefix out of range") + } + first := raw[offset] + prefix := 1 + switch { + case first < 254: + n = int(first) + case first == 254: + if len(raw)-offset < 4 { + return 0, 0, fmt.Errorf("truncated long bytes prefix") + } + n = int(raw[offset+1]) | int(raw[offset+2])<<8 | int(raw[offset+3])<<16 + prefix = 4 + default: + return 0, 0, fmt.Errorf("invalid bytes prefix") + } + total := prefix + n + padding := (4 - total%4) % 4 + if total > len(raw)-offset-padding { + return 0, 0, fmt.Errorf("truncated bytes payload") + } + return n, total + padding, nil +} diff --git a/internal/rpc/request_preflight_test.go b/internal/rpc/request_preflight_test.go new file mode 100644 index 00000000..0ca59a81 --- /dev/null +++ b/internal/rpc/request_preflight_test.go @@ -0,0 +1,134 @@ +package rpc + +import ( + "context" + "encoding/binary" + "testing" + + "github.com/gotd/td/bin" + "github.com/gotd/td/clock" + "github.com/gotd/td/tg" + "github.com/gotd/td/tgerr" + "go.uber.org/zap/zaptest" + + appfiles "telesrv/internal/app/files" +) + +func TestRequestVectorPreflightMirrorsHandlerCaps(t *testing.T) { + for id, policy := range requestVectorPolicies { + id, policy := id, policy + t.Run(tlTypeName(id), func(t *testing.T) { + atCap := fixedVectorRequest(id, policy, policy.max) + if err := preflightRPCRequest(id, &bin.Buffer{Buf: atCap}); err != nil { + t.Fatalf("cap=%d rejected: %v", policy.max, err) + } + over := fixedVectorRequest(id, policy, policy.max+1) + err := preflightRPCRequest(id, &bin.Buffer{Buf: over}) + if err == nil { + t.Fatalf("cap+1=%d accepted", policy.max+1) + } + want := "LIMIT_INVALID" + if id == tg.UsersGetUsersRequestTypeID { + want = "INPUT_REQUEST_TOO_LONG" + } + if !tgerr.Is(err, want) { + t.Fatalf("cap+1 error = %v, want %s", err, want) + } + }) + } +} + +func TestRequestVectorPreflightRejectsForgedCountInConstantSpace(t *testing.T) { + policy := requestVectorPolicies[tg.UsersGetUsersRequestTypeID] + raw := fixedVectorRequest(tg.UsersGetUsersRequestTypeID, policy, 0) + binary.LittleEndian.PutUint32(raw[policy.vectorOffset+4:], uint32(0x7fffffff)) + if err := preflightRPCRequest(tg.UsersGetUsersRequestTypeID, &bin.Buffer{Buf: raw}); !tgerr.Is(err, "INPUT_REQUEST_INVALID") { + t.Fatalf("forged count error = %v, want INPUT_REQUEST_INVALID", err) + } +} + +func TestRequestVectorPreflightRunsAfterWrapperBeforeTypedDecode(t *testing.T) { + ids := make([]tg.InputUserClass, 101) + for i := range ids { + ids[i] = &tg.InputUserSelf{} + } + wrapped := &tg.InvokeWithLayerRequest{Layer: 227, Query: &tg.UsersGetUsersRequest{ID: ids}} + var body bin.Buffer + if err := wrapped.Encode(&body); err != nil { + t.Fatalf("encode wrapper: %v", err) + } + r := New(Config{}, Deps{}, zaptest.NewLogger(t), clock.System) + _, err := r.Dispatch(context.Background(), [8]byte{1}, 1, &body) + if !tgerr.Is(err, "INPUT_REQUEST_TOO_LONG") { + t.Fatalf("wrapped oversized users.getUsers error = %v, want INPUT_REQUEST_TOO_LONG", err) + } +} + +func TestUploadPartPreflightBeforeBytesDecode(t *testing.T) { + for _, tc := range []struct { + name string + id uint32 + offset int + big bool + parts int + size int + want string + truncateBy int + }{ + {name: "small_at_cap", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: appfiles.MaxUploadPartBytes}, + {name: "small_over_cap", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: appfiles.MaxUploadPartBytes + 1, want: "FILE_PART_TOO_BIG"}, + {name: "big_at_cap", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, big: true, parts: appfiles.MaxUploadParts, size: appfiles.MaxUploadPartBytes}, + {name: "big_parts_over_cap", id: tg.UploadSaveBigFilePartRequestTypeID, offset: 20, big: true, parts: appfiles.MaxUploadParts + 1, size: 1, want: "FILE_PART_INVALID"}, + {name: "truncated", id: tg.UploadSaveFilePartRequestTypeID, offset: 16, size: 1024, truncateBy: 1, want: "INPUT_REQUEST_INVALID"}, + } { + t.Run(tc.name, func(t *testing.T) { + raw := uploadPartRequest(tc.id, tc.offset, tc.parts, tc.size) + if tc.truncateBy > 0 { + raw = raw[:len(raw)-tc.truncateBy] + } + err := preflightRPCRequest(tc.id, &bin.Buffer{Buf: raw}) + if tc.want == "" { + if err != nil { + t.Fatalf("preflight: %v", err) + } + return + } + if !tgerr.Is(err, tc.want) { + t.Fatalf("error = %v, want %s", err, tc.want) + } + }) + } +} + +func fixedVectorRequest(id uint32, policy requestVectorPolicy, count int) []byte { + raw := make([]byte, policy.vectorOffset+8+count*policy.minElemBytes) + binary.LittleEndian.PutUint32(raw[0:4], id) + binary.LittleEndian.PutUint32(raw[policy.vectorOffset:policy.vectorOffset+4], tlVectorTypeID) + binary.LittleEndian.PutUint32(raw[policy.vectorOffset+4:policy.vectorOffset+8], uint32(count)) + return raw +} + +func uploadPartRequest(id uint32, offset, parts, size int) []byte { + raw := make([]byte, offset) + binary.LittleEndian.PutUint32(raw[:4], id) + if offset == 20 { + binary.LittleEndian.PutUint32(raw[16:20], uint32(parts)) + } + prefix := 1 + if size >= 254 { + prefix = 4 + } + total := prefix + size + padding := (4 - total%4) % 4 + start := len(raw) + raw = append(raw, make([]byte, total+padding)...) + if prefix == 1 { + raw[start] = byte(size) + } else { + raw[start] = 254 + raw[start+1] = byte(size) + raw[start+2] = byte(size >> 8) + raw[start+3] = byte(size >> 16) + } + return raw +} diff --git a/internal/rpc/router.go b/internal/rpc/router.go index 24862f73..3bf6ac61 100644 --- a/internal/rpc/router.go +++ b/internal/rpc/router.go @@ -3,6 +3,7 @@ package rpc import ( "context" "encoding/hex" + "errors" "fmt" "sync" "time" @@ -53,6 +54,12 @@ type Config struct { OutboundPushTimeout time.Duration SendRateLimit int SendRateWindow time.Duration + // AuthCode*RateLimit protects the unauthenticated sendCode/resendCode write path. + // The phone budget is keyed by SHA-256(normalized phone), never by the plaintext phone; + // the second budget is keyed by the physical connection's raw auth_key_id. + AuthCodePhoneRateLimit int + AuthCodeAuthKeyRateLimit int + AuthCodeRateWindow time.Duration // CatchupRateLimit/CatchupRateWindow 限制 difference 类 catch-up RPC（getChannelDifference / // getPeerDialogs）的每用户频率（设计 Phase 2 / §10.3）：nudge 被消费后客户端会触发这两类 // catch-up，放开大群 nudge 全速前需 FLOOD_WAIT 兜底防风暴打爆 PG。两类各自独立计数、共用同一 @@ -216,6 +223,7 @@ func New(cfg Config, deps Deps, log *zap.Logger, clk clock.Clock) *Router { // 再按 TypeID 路由到 typed handler。满足 mtprotoedge.RPCHandler。 func (r *Router) Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error) { preStart := r.clock.Now() + ctx = withInboundRPCBytes(ctx, b.Len()) ctx = WithRawAuthKeyID(ctx, authKeyID) effectiveAuthKeyID, err := r.effectiveAuthKeyID(ctx, authKeyID, sessionID) if err != nil { @@ -579,8 +587,8 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En if err != nil { return nil, fmt.Errorf("decode invokeAfterMsgs msg_ids: %w", err) } - if msgIDs > maxInvokeAfterMsgIDs { - return nil, fmt.Errorf("decode invokeAfterMsgs msg_ids: too many ids %d", msgIDs) + if msgIDs < 0 || msgIDs > maxInvokeAfterMsgIDs { + return nil, fmt.Errorf("decode invokeAfterMsgs msg_ids: invalid count %d", msgIDs) } for i := 0; i < msgIDs; i++ { if _, err := b.Long(); err != nil { @@ -644,6 +652,16 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En } } } + knownRequest, structuralErr := layerwire.ValidateRoutableRequest(b.Buf) + if !knownRequest { + if structuralErr != nil { + return nil, mapLayerwirePreflightError(structuralErr) + } + // Unknown methods are opaque after the bounded/alignment check above: no generated + // decoder will touch their body. Route them directly to the compatibility fallback + // so every unknown constructor is traced, including pre-login probes. + return r.fallback(ctx, b) + } if r.deps.Auth != nil { if _, ok := UserIDFrom(ctx); !ok && !rpcAllowedWithoutAuthorization(id) { fields := append([]zap.Field{ @@ -654,6 +672,16 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En return nil, authKeyUnregisteredErr() } } + if err := preflightRPCRequest(id, b); err != nil { + return nil, err + } + // Run the same allocation-free schema walker used by layer aliases/DrKLO transforms on + // every canonical request before gotd's generated decoder materializes vectors/bytes. + // Method-specific caps above preserve exact existing RPC errors; this generic budget + // closes variable-offset/nested-object paths and malformed constructor recursion. + if structuralErr != nil { + return nil, mapLayerwirePreflightError(structuralErr) + } // 任何未包 invokeWithoutUpdates 的已登录 RPC 都把当前 session 视为 updates // 接收者。仅靠 updates.getState/getDifference 置位会漏掉 DrKLO 热恢复： // 它重连后不重建同步基线（pts 在进程内存里），只发普通业务请求，置位 @@ -682,6 +710,13 @@ func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.En } } +func mapLayerwirePreflightError(err error) error { + if errors.Is(err, layerwire.ErrResourceLimit) { + return inputRequestTooLongErr() + } + return inputRequestInvalidErr() +} + func tlTypeName(id uint32) string { tlTypeNamesOnce.Do(func() { names := tg.NamesMap() diff --git a/internal/rpc/router_auth_cache_test.go b/internal/rpc/router_auth_cache_test.go index 6443d735..1ec10459 100644 --- a/internal/rpc/router_auth_cache_test.go +++ b/internal/rpc/router_auth_cache_test.go @@ -6,6 +6,7 @@ import ( "github.com/gotd/td/bin" "github.com/gotd/td/clock" + "github.com/gotd/td/proto" "github.com/gotd/td/tg" "github.com/gotd/td/tgerr" "go.uber.org/zap/zaptest" @@ -13,13 +14,35 @@ import ( "time" ) +// authBindingCaptureSessions keeps session authorization state separate from the target of an +// asynchronous presence push. The broad captureSessions fake intentionally records the latest +// PushToUser target in userID, which is useful to most RPC tests but can race a stale-auth-key +// assertion and make an old presence echo look like the session was rebound. +type authBindingCaptureSessions struct { + *captureSessions +} + +func newAuthBindingCaptureSessions() *authBindingCaptureSessions { + return &authBindingCaptureSessions{captureSessions: &captureSessions{}} +} + +func (s *authBindingCaptureSessions) PushToUserExceptSession(_ context.Context, userID, _ int64, t proto.MessageType, msg bin.Encoder) (int, error) { + s.mu.Lock() + defer s.mu.Unlock() + s.messageType = t + s.message = msg + s.userMessage = msg + s.pushUserIDs = append(s.pushUserIDs, userID) + return 1, nil +} + func TestDispatchPromotesNegativeSessionCacheFromPositiveAuthCache(t *testing.T) { authKeyID := [8]byte{0x91, 0x91, 0x91, 0x91, 0x91, 0x91, 0x91, 0x91} const ( sessionID = int64(300) userID = int64(1000000001) ) - sessions := &captureSessions{} + sessions := newAuthBindingCaptureSessions() sessions.BindAuthKey(sessionID, authKeyID) sessions.BindUser(sessionID, 0) auth := &captureAuthService{} @@ -55,7 +78,7 @@ func TestDispatchPromotesNegativeSessionCacheFromPositiveAuthCache(t *testing.T) func TestBindTempAuthKeyClearsNegativeUserCache(t *testing.T) { var tempAuthKeyID = [8]byte{0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55, 0x55} var permAuthKeyID = [8]byte{0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11, 0x11} - sessions := &captureSessions{} + sessions := newAuthBindingCaptureSessions() auth := &captureAuthService{} r := New(Config{}, Deps{ Auth: auth, @@ -89,7 +112,7 @@ func TestBindTempAuthKeyClearsNegativeUserCache(t *testing.T) { func TestDispatchRevalidatesCachedTempAuthKeyBinding(t *testing.T) { var tempAuthKeyID = [8]byte{0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65, 0x65} var permAuthKeyID = [8]byte{0x21, 0x21, 0x21, 0x21, 0x21, 0x21, 0x21, 0x21} - sessions := &captureSessions{} + sessions := newAuthBindingCaptureSessions() auth := &captureAuthService{ resolvedAuthKeyID: permAuthKeyID, hasResolved: true, @@ -142,7 +165,7 @@ func TestDispatchRevalidatesCachedTempAuthKeyBinding(t *testing.T) { func TestDispatchUsesCachedTempAuthKeyUserUntilWriteSideInvalidation(t *testing.T) { var tempAuthKeyID = [8]byte{0x66, 0x66, 0x66, 0x66, 0x66, 0x66, 0x66, 0x66} var permAuthKeyID = [8]byte{0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22, 0x22} - sessions := &captureSessions{} + sessions := newAuthBindingCaptureSessions() auth := &captureAuthService{ resolvedAuthKeyID: permAuthKeyID, hasResolved: true, diff --git a/internal/rpc/rpc_testkit_updates_test.go b/internal/rpc/rpc_testkit_updates_test.go index 8efd5c4c..e3a363a8 100644 --- a/internal/rpc/rpc_testkit_updates_test.go +++ b/internal/rpc/rpc_testkit_updates_test.go @@ -15,8 +15,10 @@ type captureUpdates struct { cleared bool date int events []domain.UpdateEvent + excludeAuthKeyID [8]byte excludeSessionID int64 reliableDispatch bool + difference *domain.UpdateDifference } func (s *captureUpdates) UsesReliableDispatch() bool { @@ -59,6 +61,9 @@ func (s *captureUpdates) AcknowledgeCurrentState(_ context.Context, authKeyID [8 func (s *captureUpdates) GetDifference(_ context.Context, authKeyID [8]byte, userID int64, _ domain.UpdateState) (domain.UpdateDifference, error) { s.authKeyID = authKeyID s.userID = userID + if s.difference != nil { + return *s.difference, nil + } return domain.UpdateDifference{State: s.state}, nil } @@ -90,8 +95,13 @@ func (s *captureUpdates) PublishNewMessage(_ context.Context, userID int64, msg return event, st, nil } -func (s *captureUpdates) RecordStory(_ context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { +func (s *captureUpdates) captureExclude(excludeAuthKeyID [8]byte, excludeSessionID int64) { + s.excludeAuthKeyID = excludeAuthKeyID s.excludeSessionID = excludeSessionID +} + +func (s *captureUpdates) RecordStory(_ context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{ Type: domain.UpdateEventStory, Peer: story.Owner, @@ -109,10 +119,10 @@ func (s *captureUpdates) RecordStoryFanout(_ context.Context, userID int64, stor }) } -func (s *captureUpdates) RecordReadHistory(_ context.Context, authKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { +func (s *captureUpdates) RecordReadHistory(_ context.Context, authKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { s.authKeyID = authKeyID s.userID = userID - s.excludeSessionID = excludeSessionID + s.captureExclude(excludeAuthKeyID, excludeSessionID) event := domain.UpdateEvent{ Type: domain.UpdateEventReadHistoryInbox, Pts: s.state.Pts, @@ -127,8 +137,8 @@ func (s *captureUpdates) RecordReadHistory(_ context.Context, authKeyID [8]byte, return event, s.state, nil } -func (s *captureUpdates) RecordReadStories(_ context.Context, authKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordReadStories(_ context.Context, authKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{ Type: domain.UpdateEventReadStories, Peer: read.Peer, @@ -136,8 +146,8 @@ func (s *captureUpdates) RecordReadStories(_ context.Context, authKeyID [8]byte, }) } -func (s *captureUpdates) RecordSentStoryReaction(_ context.Context, authKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordSentStoryReaction(_ context.Context, authKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{ Type: domain.UpdateEventSentStoryReaction, Peer: reaction.Peer, @@ -147,8 +157,8 @@ func (s *captureUpdates) RecordSentStoryReaction(_ context.Context, authKeyID [8 }) } -func (s *captureUpdates) RecordNewStoryReaction(_ context.Context, authKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordNewStoryReaction(_ context.Context, authKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, ownerUserID, domain.UpdateEvent{ Type: domain.UpdateEventNewStoryReaction, Peer: domain.Peer{Type: domain.PeerTypeUser, ID: reaction.ViewerID}, @@ -158,8 +168,8 @@ func (s *captureUpdates) RecordNewStoryReaction(_ context.Context, authKeyID [8] }) } -func (s *captureUpdates) RecordQuickReplyMutation(_ context.Context, authKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordQuickReplyMutation(_ context.Context, authKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) event := domain.UpdateEvent{ Date: mutation.Date, QuickReplies: append([]domain.QuickReply(nil), mutation.List.QuickReplies...), @@ -183,78 +193,78 @@ func (s *captureUpdates) RecordQuickReplyMutation(_ context.Context, authKeyID [ return s.recordCapturedEvent(authKeyID, userID, event) } -func (s *captureUpdates) RecordChannelState(_ context.Context, authKeyID [8]byte, userID, channelID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordChannelState(_ context.Context, authKeyID [8]byte, userID, channelID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventChannelState, Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID}}) } -func (s *captureUpdates) RecordContactsReset(_ context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordContactsReset(_ context.Context, authKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventContactsReset}) } -func (s *captureUpdates) RecordDraftMessage(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordDraftMessage(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventDraftMessage, Peer: peer, MaxID: topMsgID}) } -func (s *captureUpdates) RecordDialogPinned(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordDialogPinned(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventDialogPinned, Peer: peer, Bool: pinned, FolderID: folderID}) } -func (s *captureUpdates) RecordPinnedDialogs(_ context.Context, authKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordPinnedDialogs(_ context.Context, authKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventPinnedDialogs, Peers: append([]domain.Peer(nil), order...), FolderID: folderID}) } -func (s *captureUpdates) RecordSavedDialogPinned(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordSavedDialogPinned(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventSavedDialogPinned, Peer: peer, Bool: pinned}) } -func (s *captureUpdates) RecordPinnedSavedDialogs(_ context.Context, authKeyID [8]byte, userID int64, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordPinnedSavedDialogs(_ context.Context, authKeyID [8]byte, userID int64, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventPinnedSavedDialogs, Peers: append([]domain.Peer(nil), order...)}) } -func (s *captureUpdates) RecordDialogUnreadMark(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordDialogUnreadMark(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventDialogUnreadMark, Peer: peer, Bool: unread}) } -func (s *captureUpdates) RecordPeerSettings(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordPeerSettings(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventPeerSettings, Peer: peer, Settings: settings}) } -func (s *captureUpdates) RecordPeerStoryBlocked(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordPeerStoryBlocked(_ context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventPeerStoryBlocked, Peer: peer, Bool: blocked}) } -func (s *captureUpdates) RecordDialogFilter(_ context.Context, authKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordDialogFilter(_ context.Context, authKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventDialogFilter, FilterID: folderID, DialogFilter: folder}) } -func (s *captureUpdates) RecordDialogFilterOrder(_ context.Context, authKeyID [8]byte, userID int64, order []int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordDialogFilterOrder(_ context.Context, authKeyID [8]byte, userID int64, order []int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventDialogFilterOrder, FilterOrder: append([]int(nil), order...)}) } -func (s *captureUpdates) RecordDialogFiltersReload(_ context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordDialogFiltersReload(_ context.Context, authKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventDialogFilters}) } -func (s *captureUpdates) RecordFolderPeers(_ context.Context, authKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordFolderPeers(_ context.Context, authKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{Type: domain.UpdateEventFolderPeers, FolderPeers: append([]domain.FolderPeerUpdate(nil), peers...)}) } -func (s *captureUpdates) RecordChannelAvailableMessages(_ context.Context, authKeyID [8]byte, userID, channelID int64, availableMinID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordChannelAvailableMessages(_ context.Context, authKeyID [8]byte, userID, channelID int64, availableMinID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{ Type: domain.UpdateEventChannelAvailable, Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID}, @@ -262,8 +272,8 @@ func (s *captureUpdates) RecordChannelAvailableMessages(_ context.Context, authK }) } -func (s *captureUpdates) RecordChannelViewForumAsMessages(_ context.Context, authKeyID [8]byte, userID, channelID int64, enabled bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordChannelViewForumAsMessages(_ context.Context, authKeyID [8]byte, userID, channelID int64, enabled bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{ Type: domain.UpdateEventChannelViewForum, Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID}, @@ -271,8 +281,8 @@ func (s *captureUpdates) RecordChannelViewForumAsMessages(_ context.Context, aut }) } -func (s *captureUpdates) RecordChannelDiscussionInbox(_ context.Context, authKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { - s.excludeSessionID = excludeSessionID +func (s *captureUpdates) RecordChannelDiscussionInbox(_ context.Context, authKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) { + s.captureExclude(excludeAuthKeyID, excludeSessionID) return s.recordCapturedEvent(authKeyID, userID, domain.UpdateEvent{ Type: domain.UpdateEventReadChannelDiscussionInbox, Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID}, diff --git a/internal/rpc/send_media.go b/internal/rpc/send_media.go index 72eeb2f8..a81dfce1 100644 --- a/internal/rpc/send_media.go +++ b/internal/rpc/send_media.go @@ -22,19 +22,21 @@ const maxContactVcardLength = 8192 // outgoingSend 是 sendOutgoing 的入参：一条已校验的出站消息。 type outgoingSend struct { - randomID int64 - message string - entities []tg.MessageEntityClass - media *domain.MessageMedia - silent bool - noforwards bool - replyToInput tg.InputReplyToClass - sendAsInput tg.InputPeerClass - replyTo *domain.MessageReply - replyToReady bool - sendAs *domain.Peer - sendAsReady bool - clearDraft bool + randomID int64 + idempotencyFingerprint []byte + idempotencyPreflighted bool + message string + entities []tg.MessageEntityClass + media *domain.MessageMedia + silent bool + noforwards bool + replyToInput tg.InputReplyToClass + sendAsInput tg.InputPeerClass + replyTo *domain.MessageReply + replyToReady bool + sendAs *domain.Peer + sendAsReady bool + clearDraft bool // replyMarkup 是 bot inline keyboard（已解析+校验；非 bot 恒 nil）。 replyMarkup *domain.MessageReplyMarkup viaBotID int64 @@ -76,24 +78,26 @@ func (r *Router) sendOutgoing(ctx context.Context, userID int64, peer domain.Pee return nil, false, err } res, err := r.deps.Channels.SendMessage(ctx, userID, domain.SendChannelMessageRequest{ - UserID: userID, - ChannelID: peer.ID, - RandomID: p.randomID, - Message: p.message, - Entities: domainMessageEntitiesForViewer(userID, p.entities), - Media: p.media, - MentionUserIDs: mentionUserIDs, - SkipRecipientLookup: true, - PostAuthor: r.channelPostAuthorName(ctx, userID), - Silent: p.silent, - NoForwards: p.noforwards, - ReplyTo: replyTo, - ViaBotID: p.viaBotID, - GroupedID: p.groupedID, - ReplyMarkup: p.replyMarkup, - RichMessage: p.richMessage, - SendAs: sendAs, - Date: int(r.clock.Now().Unix()), + UserID: userID, + ChannelID: peer.ID, + RandomID: p.randomID, + IdempotencyFingerprint: p.idempotencyFingerprint, + IdempotencyPreflighted: p.idempotencyPreflighted, + Message: p.message, + Entities: domainMessageEntitiesForViewer(userID, p.entities), + Media: p.media, + MentionUserIDs: mentionUserIDs, + SkipRecipientLookup: true, + PostAuthor: r.channelPostAuthorName(ctx, userID), + Silent: p.silent, + NoForwards: p.noforwards, + ReplyTo: replyTo, + ViaBotID: p.viaBotID, + GroupedID: p.groupedID, + ReplyMarkup: p.replyMarkup, + RichMessage: p.richMessage, + SendAs: sendAs, + Date: int(r.clock.Now().Unix()), }) if err != nil { return nil, false, channelInvalidErr(err) @@ -112,7 +116,7 @@ func (r *Router) sendOutgoing(ctx context.Context, userID int64, peer domain.Pee // 频道链接预览 pending 占位：带外解析并就地替换（异步，不阻塞发送 echo）。 r.maybeEnqueueWebPageResolve(userID, peer, res.Message.ID, res.Message.Media) } - if p.clearDraft { + if p.clearDraft && !res.Duplicate { r.clearDraftAfterSend(ctx, userID, peer, replyTo) } return updates, res.Duplicate, nil @@ -143,26 +147,30 @@ func (r *Router) sendOutgoing(ctx context.Context, userID int64, peer domain.Pee return nil, false, err } sessionID, _ := SessionIDFrom(ctx) - authKeyID, _ := AuthKeyIDFrom(ctx) + // outbox 的排除键定位的是发起 RPC 的物理连接；PFS temp key 绑定后 + // AuthKeyIDFrom 是业务视角 perm key，不能用它代替连接实际 raw key。 + authKeyID := rawAuthKeyIDForOrigin(ctx) res, err := r.deps.Messages.SendPrivateText(ctx, userID, domain.SendPrivateTextRequest{ - SenderUserID: userID, - RecipientUserID: peer.ID, - RandomID: p.randomID, - Message: p.message, - Entities: domainMessageEntitiesForViewer(userID, p.entities), - Media: p.media, - Silent: p.silent, - NoForwards: p.noforwards, - ReplyTo: replyTo, - Date: int(r.clock.Now().Unix()), - OriginAuthKeyID: authKeyID, - OriginSessionID: sessionID, - RecipientBlocked: recipientBlocked, - ReplyMarkup: p.replyMarkup, - RichMessage: p.richMessage, - ViaBotID: p.viaBotID, - GroupedID: p.groupedID, - Effect: p.effect, + SenderUserID: userID, + RecipientUserID: peer.ID, + RandomID: p.randomID, + Message: p.message, + Entities: domainMessageEntitiesForViewer(userID, p.entities), + Media: p.media, + Silent: p.silent, + NoForwards: p.noforwards, + ReplyTo: replyTo, + Date: int(r.clock.Now().Unix()), + OriginAuthKeyID: authKeyID, + OriginSessionID: sessionID, + RecipientBlocked: recipientBlocked, + IdempotencyFingerprint: p.idempotencyFingerprint, + IdempotencyPreflighted: p.idempotencyPreflighted, + ReplyMarkup: p.replyMarkup, + RichMessage: p.richMessage, + ViaBotID: p.viaBotID, + GroupedID: p.groupedID, + Effect: p.effect, }) if err != nil { fields := append(r.contextLogFields(ctx), @@ -179,9 +187,13 @@ func (r *Router) sendOutgoing(ctx context.Context, userID int64, peer domain.Pee r.log.Warn("messages.sendMessage private store failed", fields...) return nil, false, messageSendErr(err) } - users := r.usersForMessageUpdate(ctx, userID, res.SenderMessage) - chats := r.chatsForMessageUpdate(ctx, userID, res.SenderMessage) - if p.clearDraft { + var users []tg.UserClass + var chats []tg.ChatClass + if !res.Duplicate { + users = r.usersForMessageUpdate(ctx, userID, res.SenderMessage) + chats = r.chatsForMessageUpdate(ctx, userID, res.SenderMessage) + } + if p.clearDraft && !res.Duplicate { r.clearDraftAfterSend(ctx, userID, peer, replyTo) } if !res.Duplicate { @@ -189,7 +201,7 @@ func (r *Router) sendOutgoing(ctx context.Context, userID int64, peer domain.Pee r.maybeEnqueueWebPageResolve(userID, peer, res.SenderMessage.ID, res.SenderMessage.Media) r.enqueueBotAPIPrivateMessageUpdateAsync(ctx, res) } - return tgPrivateMessageUpdates(res.SenderEvent, res.SenderMessage, p.randomID, true, users, chats), res.Duplicate, nil + return tgPrivateSendResultUpdates(res, p.randomID, true, users, chats), res.Duplicate, nil } // onMessagesUploadMedia 解析 InputMedia（上传或引用），返回可复用的 tg.MessageMedia。 @@ -244,6 +256,10 @@ func (r *Router) onMessagesSendMedia(ctx context.Context, req *tg.MessagesSendMe case *tg.InputMediaEmpty, *tg.InputMediaWebPage: return r.onMessagesSendMessage(ctx, sendMessageRequestFromSendMedia(req)) } + // 指纹是幂等校验元数据，不能让一个本应由 media resolver 映射成 + // MEDIA_INVALID 的畸形 input 因 Encode 失败提前变成 INTERNAL。合法请求会得到 + // 原始 TL 指纹；编码失败则留空，由 store 使用 domain fallback。 + idempotencyFingerprint, _ := sendMediaIdempotencyFingerprint(req) // 媒体 caption 里的链接/@mention/#hashtag 等同样补自动高亮实体（客户端未带时）。 req.Entities = augmentAutoEntities(req.Message, req.Entities) userID, _, err := r.currentUserID(ctx) @@ -253,6 +269,17 @@ func (r *Router) onMessagesSendMedia(ctx context.Context, req *tg.MessagesSendMe if userID == 0 { return nil, peerIDInvalidErr() } + peer, ok := r.domainPeerFromInputPeer(userID, req.Peer) + if !ok || peer.ID == 0 { + return nil, peerIDInvalidErr() + } + replay, err := r.lookupOutgoingReplay(ctx, userID, peer, req.RandomID, idempotencyFingerprint) + if err != nil { + return nil, err + } + if replay.found { + return r.outgoingReplayUpdates(ctx, userID, peer, req.RandomID, replay), nil + } // 消息特效：仅接受 catalog 内的合法 effect id（非法 id → EFFECT_ID_INVALID，官方行为）。 if r.messageEffectInvalid(ctx, req.Effect) { return nil, effectIDInvalidErr() @@ -260,7 +287,7 @@ func (r *Router) onMessagesSendMedia(ctx context.Context, req *tg.MessagesSendMe if err := r.checkSendRateLimit(ctx, userID, 1); err != nil { return nil, err } - peer, err := r.checkedDomainPeerFromInputPeer(ctx, userID, req.Peer) + peer, err = r.checkedDomainPeerFromInputPeer(ctx, userID, req.Peer) if err != nil { return nil, err } @@ -281,29 +308,33 @@ func (r *Router) onMessagesSendMedia(ctx context.Context, req *tg.MessagesSendMe } if req.ScheduleDate != 0 && !scheduleDateIsImmediate(req.ScheduleDate, int(r.clock.Now().Unix())) { return r.scheduleOutgoing(ctx, userID, peer, outgoingSend{ - randomID: req.RandomID, - message: req.Message, - entities: req.Entities, - media: media, - silent: req.Silent, - noforwards: req.Noforwards, - replyToInput: req.ReplyTo, - sendAsInput: req.SendAs, - clearDraft: req.ClearDraft, + randomID: req.RandomID, + idempotencyFingerprint: idempotencyFingerprint, + idempotencyPreflighted: replay.checked, + message: req.Message, + entities: req.Entities, + media: media, + silent: req.Silent, + noforwards: req.Noforwards, + replyToInput: req.ReplyTo, + sendAsInput: req.SendAs, + clearDraft: req.ClearDraft, }, req.ScheduleDate, req.ScheduleRepeatPeriod) } updates, _, err := r.sendOutgoing(ctx, userID, peer, outgoingSend{ - randomID: req.RandomID, - message: req.Message, - entities: req.Entities, - media: media, - silent: req.Silent, - noforwards: req.Noforwards, - replyToInput: req.ReplyTo, - sendAsInput: req.SendAs, - clearDraft: req.ClearDraft, - replyMarkup: replyMarkup, - effect: req.Effect, + randomID: req.RandomID, + idempotencyFingerprint: idempotencyFingerprint, + idempotencyPreflighted: replay.checked, + message: req.Message, + entities: req.Entities, + media: media, + silent: req.Silent, + noforwards: req.Noforwards, + replyToInput: req.ReplyTo, + sendAsInput: req.SendAs, + clearDraft: req.ClearDraft, + replyMarkup: replyMarkup, + effect: req.Effect, }) if err != nil { return nil, err @@ -311,7 +342,7 @@ func (r *Router) onMessagesSendMedia(ctx context.Context, req *tg.MessagesSendMe return updates, nil } -// onMessagesSendMultiMedia 发送相册（多条媒体）。本阶段不绑定 grouped_id（各条作为独立消息呈现）。 +// onMessagesSendMultiMedia 发送相册（多条媒体），并在解析媒体前持久预留 grouped_id。 func (r *Router) onMessagesSendMultiMedia(ctx context.Context, req *tg.MessagesSendMultiMediaRequest) (tg.UpdatesClass, error) { userID, _, err := r.currentUserID(ctx) if err != nil { @@ -323,14 +354,20 @@ func (r *Router) onMessagesSendMultiMedia(ctx context.Context, req *tg.MessagesS if len(req.MultiMedia) == 0 || len(req.MultiMedia) > maxSendMultiMediaItems { return nil, limitInvalidErr() } - peer, err := r.checkedDomainPeerFromInputPeer(ctx, userID, req.Peer) - if err != nil { - return nil, err + peer, ok := r.domainPeerFromInputPeer(userID, req.Peer) + if !ok || peer.ID == 0 { + return nil, peerIDInvalidErr() } + randomIDs := make(map[int64]struct{}, len(req.MultiMedia)) + reservationItems := make([]domain.AlbumGroupReservationItem, 0, len(req.MultiMedia)) for _, item := range req.MultiMedia { if item.RandomID == 0 { return nil, randomIDEmptyErr() } + if _, duplicate := randomIDs[item.RandomID]; duplicate { + return nil, randomIDDuplicateErr() + } + randomIDs[item.RandomID] = struct{}{} if utf8.RuneCountInString(item.Message) > maxSendMessageTextLength { return nil, mediaCaptionTooLongErr() } @@ -340,18 +377,69 @@ func (r *Router) onMessagesSendMultiMedia(ctx context.Context, req *tg.MessagesS if item.Media == nil { return nil, mediaInvalidErr() } + intentHash, fingerprintErr := sendMultiMediaItemIdempotencyFingerprint(req, item) + if fingerprintErr != nil { + // 畸形 InputMedia 的 TL 编码失败不能变成 INTERNAL；保持 media 输入错误语义。 + return nil, mediaInvalidErr() + } + reservationItems = append(reservationItems, domain.AlbumGroupReservationItem{ + RandomID: item.RandomID, + IntentHash: intentHash, + }) } - if err := r.checkSendRateLimit(ctx, userID, len(req.MultiMedia)); err != nil { + replays := make([]outgoingReplayLookup, len(req.MultiMedia)) + absentCount := 0 + for i, item := range req.MultiMedia { + replay, err := r.lookupOutgoingReplay(ctx, userID, peer, item.RandomID, reservationItems[i].IntentHash) + if err != nil { + return nil, err + } + replays[i] = replay + if !replay.found { + absentCount++ + } + } + if absentCount == 0 { + results := make([]tg.UpdatesClass, 0, len(req.MultiMedia)) + for i, item := range req.MultiMedia { + results = append(results, r.outgoingReplayUpdates(ctx, userID, peer, item.RandomID, replays[i])) + } + return combineSendUpdates(results), nil + } + if err := r.checkSendRateLimit(ctx, userID, absentCount); err != nil { + return nil, err + } + peer, err = r.checkedDomainPeerFromInputPeer(ctx, userID, req.Peer) + if err != nil { return nil, err } - combined := make([]tg.UpdateClass, 0, len(req.MultiMedia)*2) - usersByID := map[int64]tg.UserClass{} - chatsByID := map[int64]tg.ChatClass{} - date := 0 - // 整个 album 共享一个 grouped_id，客户端据此把各条渲染成一个相册组。 - groupedID := randomNonZeroInt64() + // 必须在 resolveInputMedia 或发送任何 item 之前原子预留：首次请求若在第 N 条 + // 失败，客户端只重试失败子集时仍从已绑定 random_id 恢复整包 grouped_id。 + groupedID, err := r.reserveAlbumGroup(ctx, userID, peer, reservationItems) + if err != nil { + return nil, err + } + for _, replay := range replays { + if replay.found { + messageGroupedID := replay.private.SenderMessage.GroupedID + if peer.Type == domain.PeerTypeChannel { + messageGroupedID = replay.channel.Message.GroupedID + } + if messageGroupedID != groupedID { + return nil, internalErr() + } + } + } + + results := make([]tg.UpdatesClass, 0, len(req.MultiMedia)) + clearDraftPending := req.ClearDraft for i, item := range req.MultiMedia { + idempotencyFingerprint := reservationItems[i].IntentHash + if replays[i].found { + results = append(results, r.outgoingReplayUpdates(ctx, userID, peer, item.RandomID, replays[i])) + continue + } media, err := r.resolveInputMedia(ctx, userID, item.Media) if err != nil { return nil, err @@ -360,49 +448,35 @@ func (r *Router) onMessagesSendMultiMedia(ctx context.Context, req *tg.MessagesS return nil, mediaInvalidErr() } p := outgoingSend{ - randomID: item.RandomID, - message: item.Message, - entities: augmentAutoEntities(item.Message, item.Entities), - media: media, - silent: req.Silent, - noforwards: req.Noforwards, - replyToInput: req.ReplyTo, - sendAsInput: req.SendAs, - clearDraft: req.ClearDraft && i == 0, - groupedID: groupedID, + randomID: item.RandomID, + idempotencyFingerprint: idempotencyFingerprint, + idempotencyPreflighted: replays[i].checked, + message: item.Message, + entities: augmentAutoEntities(item.Message, item.Entities), + media: media, + silent: req.Silent, + noforwards: req.Noforwards, + replyToInput: req.ReplyTo, + sendAsInput: req.SendAs, + clearDraft: clearDraftPending, + groupedID: groupedID, } var result tg.UpdatesClass + duplicate := false if req.ScheduleDate != 0 && !scheduleDateIsImmediate(req.ScheduleDate, int(r.clock.Now().Unix())) { result, err = r.scheduleOutgoing(ctx, userID, peer, p, req.ScheduleDate, 0) } else { - result, _, err = r.sendOutgoing(ctx, userID, peer, p) + result, duplicate, err = r.sendOutgoing(ctx, userID, peer, p) } if err != nil { return nil, err } - if upd, ok := result.(*tg.Updates); ok { - combined = append(combined, upd.Updates...) - for _, u := range upd.Users { - if id := userClassID(u); id != 0 { - usersByID[id] = u - } - } - for _, c := range upd.Chats { - if id := chatClassID(c); id != 0 { - chatsByID[id] = c - } - } - if upd.Date != 0 { - date = upd.Date - } + if p.clearDraft && !duplicate { + clearDraftPending = false } + results = append(results, result) } - return &tg.Updates{ - Updates: combined, - Users: mapValuesUsers(usersByID), - Chats: mapValuesChats(chatsByID), - Date: date, - }, nil + return combineSendUpdates(results), nil } // resolveInputMedia 把 tg.InputMedia 解析为 domain.MessageMedia（上传则落库，引用则加载）。 diff --git a/internal/rpc/send_media_test.go b/internal/rpc/send_media_test.go index d2dee278..0e94b85a 100644 --- a/internal/rpc/send_media_test.go +++ b/internal/rpc/send_media_test.go @@ -836,6 +836,126 @@ func TestSendMediaPrivateSticker(t *testing.T) { } } +func TestSendMultiMediaPartialFailureSubsetRetryKeepsReservedGroupedID(t *testing.T) { + ctx := context.Background() + r, owner, friend := newMediaTestRouter(t) + files := r.deps.Files.(*fakeFiles) + + first := tg.InputSingleMedia{ + Media: &tg.InputMediaDocument{ID: &tg.InputDocument{ID: 555, AccessHash: 5}}, + RandomID: 41001, + Message: "first", + } + second := tg.InputSingleMedia{ + // 首次请求时 556 尚不存在，使第一条已提交后第二条解析失败。 + Media: &tg.InputMediaDocument{ID: &tg.InputDocument{ID: 556, AccessHash: 6}}, + RandomID: 41002, + Message: "second", + } + peer := &tg.InputPeerUser{UserID: friend.ID, AccessHash: friend.AccessHash} + if _, err := r.onMessagesSendMultiMedia(WithUserID(ctx, owner.ID), &tg.MessagesSendMultiMediaRequest{ + Peer: peer, + MultiMedia: []tg.InputSingleMedia{first, second}, + }); err == nil || !tgerr.Is(err, "MEDIA_INVALID") { + t.Fatalf("partial album err=%v, want MEDIA_INVALID after first item commit", err) + } + + files.docs[556] = domain.Document{ID: 556, AccessHash: 6, DCID: 2, MimeType: "image/jpeg"} + retry, err := r.onMessagesSendMultiMedia(WithUserID(ctx, owner.ID), &tg.MessagesSendMultiMediaRequest{ + Peer: peer, + MultiMedia: []tg.InputSingleMedia{second}, + }) + if err != nil { + t.Fatalf("retry failed subset: %v", err) + } + retryMessage := newMessageFromUpdates(t, retry) + retryGroup, ok := retryMessage.GetGroupedID() + if !ok || retryGroup == 0 { + t.Fatalf("retry grouped_id = %d present=%v, want non-zero reservation", retryGroup, ok) + } + + history, err := r.deps.Messages.GetHistory(ctx, owner.ID, domain.MessageFilter{ + HasPeer: true, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: friend.ID}, + Limit: 10, + }) + if err != nil { + t.Fatalf("album history: %v", err) + } + groups := make(map[int64]int64, 2) + for _, message := range history.Messages { + if message.RandomID == first.RandomID || message.RandomID == second.RandomID { + groups[message.RandomID] = message.GroupedID + } + } + if len(groups) != 2 || groups[first.RandomID] != retryGroup || groups[second.RandomID] != retryGroup { + t.Fatalf("history album groups=%v, want both %d", groups, retryGroup) + } + + changed := second + changed.Message = "changed durable intent" + if _, err := r.onMessagesSendMultiMedia(WithUserID(ctx, owner.ID), &tg.MessagesSendMultiMediaRequest{ + Peer: peer, + MultiMedia: []tg.InputSingleMedia{changed}, + }); err == nil || !tgerr.Is(err, "RANDOM_ID_DUPLICATE") { + t.Fatalf("changed reserved item err=%v, want RANDOM_ID_DUPLICATE", err) + } +} + +func TestSendMultiMediaChannelPartialFailureSubsetRetryKeepsReservedGroupedID(t *testing.T) { + f := newRPCChannelFixture(t) + owner := f.user(51, "15550009401", "AlbumOwner") + member := f.user(52, "15550009402", "AlbumMember") + channel := f.createLegacyMegagroup(owner, "Album Group", member) + messageStore := memory.NewMessageStore() + f.router.deps.Messages = appmessages.NewService(messageStore, nil) + files := &fakeFiles{docs: map[int64]domain.Document{ + 555: {ID: 555, AccessHash: 5, DCID: 2, MimeType: "image/jpeg"}, + }, photos: map[int64]domain.Photo{}} + f.router.deps.Files = files + + first := tg.InputSingleMedia{ + Media: &tg.InputMediaDocument{ID: &tg.InputDocument{ID: 555, AccessHash: 5}}, RandomID: 42001, Message: "first", + } + second := tg.InputSingleMedia{ + Media: &tg.InputMediaDocument{ID: &tg.InputDocument{ID: 556, AccessHash: 6}}, RandomID: 42002, Message: "second", + } + peer := inputPeerChannel(channel) + if _, err := f.router.onMessagesSendMultiMedia(f.userCtx(owner), &tg.MessagesSendMultiMediaRequest{ + Peer: peer, MultiMedia: []tg.InputSingleMedia{first, second}, + }); err == nil || !tgerr.Is(err, "MEDIA_INVALID") { + t.Fatalf("partial channel album err=%v, want MEDIA_INVALID", err) + } + files.docs[556] = domain.Document{ID: 556, AccessHash: 6, DCID: 2, MimeType: "image/jpeg"} + retry, err := f.router.onMessagesSendMultiMedia(f.userCtx(owner), &tg.MessagesSendMultiMediaRequest{ + Peer: peer, MultiMedia: []tg.InputSingleMedia{second}, + }) + if err != nil { + t.Fatalf("retry channel subset: %v", err) + } + retryMessage := newMessageFromUpdates(t, retry) + retryGroup, ok := retryMessage.GetGroupedID() + if !ok || retryGroup == 0 { + t.Fatalf("channel retry grouped_id=%d present=%v, want non-zero", retryGroup, ok) + } + history, err := f.router.deps.Channels.GetHistory(f.ctx, owner.ID, domain.ChannelHistoryFilter{ + ChannelID: channel.ID, + Limit: 10, + }) + if err != nil { + t.Fatalf("channel album history: %v", err) + } + groups := make(map[int64]int64, 2) + for _, message := range history.Messages { + if message.RandomID == first.RandomID || message.RandomID == second.RandomID { + groups[message.RandomID] = message.GroupedID + } + } + if len(groups) != 2 || groups[first.RandomID] != retryGroup || groups[second.RandomID] != retryGroup { + t.Fatalf("channel album groups=%v, want both %d", groups, retryGroup) + } +} + func TestTGMessageMediaDocumentMarksHistoricalStickerNopremium(t *testing.T) { media := tgMessageMedia(&domain.MessageMedia{ Kind: domain.MessageMediaKindDocument, diff --git a/internal/rpc/send_replay.go b/internal/rpc/send_replay.go new file mode 100644 index 00000000..fc220bbe --- /dev/null +++ b/internal/rpc/send_replay.go @@ -0,0 +1,118 @@ +package rpc + +import ( + "context" + "crypto/sha256" + + "github.com/gotd/td/tg" + + "telesrv/internal/domain" +) + +// These optional capabilities keep the broad RPC service interfaces stable for compatibility +// fakes while production app services expose the read-only receipt lookup. +type privateSendReplayService interface { + LookupPrivateSendReplay(ctx context.Context, userID int64, req domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) +} + +type channelSendReplayService interface { + LookupChannelSendReplay(ctx context.Context, userID int64, req domain.ChannelSendReplayRequest) (domain.SendChannelMessageResult, bool, error) +} + +type outgoingReplayLookup struct { + private domain.SendPrivateTextResult + channel domain.SendChannelMessageResult + found bool + checked bool +} + +// lookupOutgoingReplay is deliberately limited to authenticated sender + destination scope + +// immutable fingerprint. It does not resolve media/replies/send-as/source messages, consume rate +// budget, run a send permission gate or emit any realtime/durable side effect. +func (r *Router) lookupOutgoingReplay(ctx context.Context, userID int64, peer domain.Peer, randomID int64, fingerprint []byte) (outgoingReplayLookup, error) { + if randomID == 0 || len(fingerprint) != sha256.Size { + return outgoingReplayLookup{}, nil + } + switch peer.Type { + case domain.PeerTypeUser: + service, ok := r.deps.Messages.(privateSendReplayService) + if !ok { + return outgoingReplayLookup{}, nil + } + res, found, err := service.LookupPrivateSendReplay(ctx, userID, domain.PrivateSendReplayRequest{ + SenderUserID: userID, + RecipientUserID: peer.ID, + RandomID: randomID, + IdempotencyFingerprint: fingerprint, + }) + if err != nil { + return outgoingReplayLookup{checked: true}, messageSendErr(err) + } + return outgoingReplayLookup{private: res, found: found, checked: true}, nil + case domain.PeerTypeChannel: + return r.lookupChannelSendReplay(ctx, userID, peer.ID, domain.Peer{}, randomID, fingerprint) + default: + return outgoingReplayLookup{}, peerIDInvalidErr() + } +} + +func (r *Router) lookupChannelSendReplay(ctx context.Context, userID, channelID int64, savedPeer domain.Peer, randomID int64, fingerprint []byte) (outgoingReplayLookup, error) { + if randomID == 0 || len(fingerprint) != sha256.Size { + return outgoingReplayLookup{}, nil + } + service, ok := r.deps.Channels.(channelSendReplayService) + if !ok { + return outgoingReplayLookup{}, nil + } + res, found, err := service.LookupChannelSendReplay(ctx, userID, domain.ChannelSendReplayRequest{ + ChannelID: channelID, + SenderUserID: userID, + SavedPeer: savedPeer, + RandomID: randomID, + IdempotencyFingerprint: fingerprint, + }) + if err != nil { + return outgoingReplayLookup{checked: true}, channelInvalidErr(err) + } + return outgoingReplayLookup{channel: res, found: found, checked: true}, nil +} + +func (r *Router) outgoingReplayUpdates(ctx context.Context, userID int64, peer domain.Peer, randomID int64, replay outgoingReplayLookup) tg.UpdatesClass { + if peer.Type == domain.PeerTypeChannel { + return r.channelMessageUpdatesWithPeerCache(ctx, userID, replay.channel, randomID, newViewerPeerCache(r)) + } + return tgPrivateSendResultUpdates(replay.private, randomID, true, nil, nil) +} + +func combineSendUpdates(results []tg.UpdatesClass) *tg.Updates { + combined := make([]tg.UpdateClass, 0, len(results)*2) + usersByID := map[int64]tg.UserClass{} + chatsByID := map[int64]tg.ChatClass{} + date := 0 + for _, result := range results { + upd, ok := result.(*tg.Updates) + if !ok || upd == nil { + continue + } + combined = append(combined, upd.Updates...) + for _, user := range upd.Users { + if id := userClassID(user); id != 0 { + usersByID[id] = user + } + } + for _, chat := range upd.Chats { + if id := chatClassID(chat); id != 0 { + chatsByID[id] = chat + } + } + if upd.Date > date { + date = upd.Date + } + } + return &tg.Updates{ + Updates: combined, + Users: mapValuesUsers(usersByID), + Chats: mapValuesChats(chatsByID), + Date: date, + } +} diff --git a/internal/rpc/stories.go b/internal/rpc/stories.go index 3e6886c4..15516892 100644 --- a/internal/rpc/stories.go +++ b/internal/rpc/stories.go @@ -2372,7 +2372,7 @@ func (r *Router) onStoriesReadStories(ctx context.Context, req *tg.StoriesReadSt if read.Advanced && r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - if _, _, err := r.deps.Updates.RecordReadStories(ctx, authKeyID, userID, read, sessionID); err != nil { + if _, _, err := r.deps.Updates.RecordReadStories(ctx, authKeyID, userID, read, rawAuthKeyIDForOrigin(ctx), sessionID); err != nil { return nil, internalErr() } } @@ -2689,11 +2689,11 @@ func (r *Router) onStoriesSendReaction(ctx context.Context, req *tg.StoriesSendR if res.Changed && r.deps.Updates != nil { authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - if _, _, err := r.deps.Updates.RecordSentStoryReaction(ctx, authKeyID, userID, res, sessionID); err != nil { + if _, _, err := r.deps.Updates.RecordSentStoryReaction(ctx, authKeyID, userID, res, rawAuthKeyIDForOrigin(ctx), sessionID); err != nil { return nil, internalErr() } if ownerUserID, ok := ownerStoryReactionNotificationUserID(res, userID); ok && res.Reaction != nil { - event, _, err := r.deps.Updates.RecordNewStoryReaction(ctx, [8]byte{}, ownerUserID, res, 0) + event, _, err := r.deps.Updates.RecordNewStoryReaction(ctx, [8]byte{}, ownerUserID, res, [8]byte{}, 0) if err != nil { return nil, internalErr() } @@ -3028,7 +3028,7 @@ func (r *Router) recordStoryChange(ctx context.Context, userID int64, story doma } authKeyID, _ := AuthKeyIDFrom(ctx) sessionID, _ := SessionIDFrom(ctx) - if _, _, err := r.deps.Updates.RecordStory(ctx, authKeyID, userID, story, sessionID); err != nil { + if _, _, err := r.deps.Updates.RecordStory(ctx, authKeyID, userID, story, rawAuthKeyIDForOrigin(ctx), sessionID); err != nil { return internalErr() } return nil diff --git a/internal/rpc/temp_key_cache.go b/internal/rpc/temp_key_cache.go index cb4b303a..7767846c 100644 --- a/internal/rpc/temp_key_cache.go +++ b/internal/rpc/temp_key_cache.go @@ -6,7 +6,7 @@ import ( "time" ) -const defaultTempKeyResolveCacheMaxEntries = 4096 +const defaultTempKeyResolveCacheMaxEntries = 262144 type tempKeyResolveEntry struct { perm [8]byte @@ -22,6 +22,7 @@ type tempKeyResolveCache struct { mu sync.Mutex max int entries map[[8]byte]*list.Element + byPerm map[[8]byte]map[[8]byte]struct{} order *list.List } @@ -30,8 +31,12 @@ func newTempKeyResolveCache(maxEntries int) *tempKeyResolveCache { maxEntries = defaultTempKeyResolveCacheMaxEntries } return &tempKeyResolveCache{ - max: maxEntries, - entries: make(map[[8]byte]*list.Element, maxEntries), + max: maxEntries, + // maxEntries is an eviction ceiling, not an expected steady-state population. A + // capacity hint of 262k eagerly reserves a large hash table for every Router even when + // PFS/temp keys are never used; let the map grow lazily with actual bindings instead. + entries: make(map[[8]byte]*list.Element), + byPerm: make(map[[8]byte]map[[8]byte]struct{}), order: list.New(), } } @@ -55,20 +60,25 @@ func (c *tempKeyResolveCache) Get(rawAuthKeyID, expectedPermAuthKeyID [8]byte, n return item.entry.perm, true } -func (c *tempKeyResolveCache) Store(rawAuthKeyID, permAuthKeyID [8]byte, expireAt, now time.Time) { +func (c *tempKeyResolveCache) Store(rawAuthKeyID, permAuthKeyID [8]byte, expireAt, _ time.Time) { if c == nil || c.max <= 0 || rawAuthKeyID == ([8]byte{}) || permAuthKeyID == ([8]byte{}) { return } c.mu.Lock() defer c.mu.Unlock() if el := c.entries[rawAuthKeyID]; el != nil { + old := el.Value.(tempKeyResolveCacheItem) + if old.entry.perm != permAuthKeyID { + c.removeReverseLocked(old.raw, old.entry.perm) + c.addReverseLocked(rawAuthKeyID, permAuthKeyID) + } el.Value = tempKeyResolveCacheItem{raw: rawAuthKeyID, entry: tempKeyResolveEntry{perm: permAuthKeyID, expireAt: expireAt}} c.order.MoveToBack(el) return } - c.evictExpiredLocked(now) el := c.order.PushBack(tempKeyResolveCacheItem{raw: rawAuthKeyID, entry: tempKeyResolveEntry{perm: permAuthKeyID, expireAt: expireAt}}) c.entries[rawAuthKeyID] = el + c.addReverseLocked(rawAuthKeyID, permAuthKeyID) for len(c.entries) > c.max { c.removeElementLocked(c.order.Front()) } @@ -91,35 +101,43 @@ func (c *tempKeyResolveCache) DeleteByPerm(permAuthKeyID [8]byte) [][8]byte { } c.mu.Lock() defer c.mu.Unlock() - rawAuthKeyIDs := make([][8]byte, 0) - for el := c.order.Front(); el != nil; { - next := el.Next() - item := el.Value.(tempKeyResolveCacheItem) - if item.entry.perm == permAuthKeyID { - rawAuthKeyIDs = append(rawAuthKeyIDs, item.raw) + raws := c.byPerm[permAuthKeyID] + rawAuthKeyIDs := make([][8]byte, 0, len(raws)) + for raw := range raws { + rawAuthKeyIDs = append(rawAuthKeyIDs, raw) + if el := c.entries[raw]; el != nil { c.removeElementLocked(el) } - el = next } return rawAuthKeyIDs } -func (c *tempKeyResolveCache) evictExpiredLocked(now time.Time) { - for el := c.order.Front(); el != nil; { - next := el.Next() - item := el.Value.(tempKeyResolveCacheItem) - if !item.entry.expireAt.After(now) { - c.removeElementLocked(el) - } - el = next - } -} - func (c *tempKeyResolveCache) removeElementLocked(el *list.Element) { if el == nil { return } item := el.Value.(tempKeyResolveCacheItem) delete(c.entries, item.raw) + c.removeReverseLocked(item.raw, item.entry.perm) c.order.Remove(el) } + +func (c *tempKeyResolveCache) addReverseLocked(raw, perm [8]byte) { + raws := c.byPerm[perm] + if raws == nil { + raws = make(map[[8]byte]struct{}) + c.byPerm[perm] = raws + } + raws[raw] = struct{}{} +} + +func (c *tempKeyResolveCache) removeReverseLocked(raw, perm [8]byte) { + raws := c.byPerm[perm] + if raws == nil { + return + } + delete(raws, raw) + if len(raws) == 0 { + delete(c.byPerm, perm) + } +} diff --git a/internal/rpc/updates.go b/internal/rpc/updates.go index b9c18530..a5068f8d 100644 --- a/internal/rpc/updates.go +++ b/internal/rpc/updates.go @@ -14,11 +14,11 @@ func (r *Router) registerUpdates(d *tg.ServerDispatcher) { d.OnUpdatesGetDifference(r.onUpdatesGetDifference) } -// onUpdatesGetState 处理 updates.getState：返回账号当前最新连续状态并推进该设备 -// 的确认水位。协议语义是客户端宣告「从现在开始同步」，启动期离线数据由 -// getDialogs 快照承载——返回设备旧确认水位会让 TDesktop（不持久化 pts、每次 -// 启动都调 getState）在 getDialogs 快照之上重放历史差分，未读重复累计、 -// dialog 预览被旧消息抢占。 +// onUpdatesGetState 处理 updates.getState。TDesktop 与 DrKLO 的启动路径把它当作 +// 「从当前快照开始同步」的显式 baseline：返回账号当前连续水位并推进该设备 observed。 +// 对无法识别的客户端仍返回同一 current state，但不把尚未被客户端带回的服务端快照 +// 记成 observed；这保留 durable difference tail，避免把 TDesktop/DrKLO 的兼容例外 +// 扩散成所有客户端都能跨过未实际确认事件的 retention 后门。 func (r *Router) onUpdatesGetState(ctx context.Context) (*tg.UpdatesState, error) { id, _ := AuthKeyIDFrom(ctx) userID, _, err := r.currentUserID(ctx) @@ -29,7 +29,16 @@ func (r *Router) onUpdatesGetState(ctx context.Context) (*tg.UpdatesState, error r.markSessionReceivesUpdates(ctx, userID) return &tg.UpdatesState{Date: int(r.clock.Now().Unix()), Qts: r.deviceEncryptedQts(ctx)}, nil } - st, err := r.deps.Updates.AcknowledgeCurrentState(ctx, id, userID) + var st domain.UpdateState + if getStateEstablishesObservedBaseline(ctx) { + st, err = r.deps.Updates.AcknowledgeCurrentState(ctx, id, userID) + } else { + st, err = r.deps.Updates.CurrentState(ctx, userID) + if err == nil { + r.log.Warn("updates.getState returned current snapshot without advancing observed baseline for unknown client", + r.contextLogFields(ctx)...) + } + } if err != nil { return nil, internalErr() } @@ -41,6 +50,15 @@ func (r *Router) onUpdatesGetState(ctx context.Context) (*tg.UpdatesState, error return ptr(out), nil } +func getStateEstablishesObservedBaseline(ctx context.Context) bool { + switch ClientTypeFrom(ctx) { + case ClientTypeTDesktop, ClientTypeAndroid: + return true + default: + return false + } +} + func (r *Router) onUpdatesGetDifference(ctx context.Context, req *tg.UpdatesGetDifferenceRequest) (tg.UpdatesDifferenceClass, error) { id, _ := AuthKeyIDFrom(ctx) userID, _, err := r.currentUserID(ctx) @@ -76,7 +94,7 @@ func (r *Router) onUpdatesGetDifference(ctx context.Context, req *tg.UpdatesGetD encMsgs, newQts := r.encryptedDifference(ctx, req.Qts) // 密聊握手/已读状态事件（无 qts）：按未投递标记补回 OtherUpdates。 stateUpdates, statePeerUserIDs, stateEventIDs := r.encryptedStateUpdates(ctx, userID) - if len(st.Events) == 0 && len(st.ChannelNudges) == 0 && len(encMsgs) == 0 && len(stateUpdates) == 0 { + if !st.Partial && len(st.Events) == 0 && len(st.ChannelNudges) == 0 && len(encMsgs) == 0 && len(stateUpdates) == 0 { r.registerBootstrapAfterBaseline(ctx, userID) return &tg.UpdatesDifferenceEmpty{Date: st.State.Date, Seq: st.State.Seq}, nil } diff --git a/internal/rpc/updates_rpc_test.go b/internal/rpc/updates_rpc_test.go index 42a69c0d..2df43805 100644 --- a/internal/rpc/updates_rpc_test.go +++ b/internal/rpc/updates_rpc_test.go @@ -17,7 +17,7 @@ import ( "telesrv/internal/store/memory" ) -func TestBootstrapLoginMessagePublishesNewMessageAfterReady(t *testing.T) { +func TestSignUpBootstrapLoginMessagePublishesNewMessageAfterReady(t *testing.T) { bootstrap := memory.NewBootstrapUpdateJobStore() updates := &captureUpdates{state: domain.UpdateState{Pts: 3, Date: 1700000000}} messages := &captureMessages{ @@ -70,7 +70,44 @@ func TestBootstrapLoginMessagePublishesNewMessageAfterReady(t *testing.T) { } } -func TestUpdatesGetStatePublishesBootstrapAfterRPCResult(t *testing.T) { +func TestSignUpBootstrapPendingJobFollowsSameAuthKeyReconnect(t *testing.T) { + bootstrap := memory.NewBootstrapUpdateJobStore() + updates := &captureUpdates{state: domain.UpdateState{Pts: 0, Date: 1700000000}} + msg := domain.Message{ + ID: 7, OwnerUserID: 1780243777, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID}, + Date: 1700000100, Body: "Login code: 12345", + } + r := New(Config{}, Deps{ + BootstrapUpdates: bootstrap, + Updates: updates, + Messages: &captureMessages{list: domain.MessageList{Messages: []domain.Message{msg}}}, + }, zaptest.NewLogger(t), clock.System) + authKeyID := [8]byte{4, 5, 6} + oldSessionID := int64(1001) + newSessionID := int64(2002) + r.enqueueLoginMessageBootstrap( + WithSessionID(WithAuthKeyID(context.Background(), authKeyID), oldSessionID), + msg, + ) + + if ready, err := bootstrap.MarkReadyForSession(context.Background(), msg.OwnerUserID, [8]byte{9}, newSessionID); err != nil || ready != 0 { + t.Fatalf("different auth-key ready=%d err=%v, want 0/nil", ready, err) + } + ready, err := bootstrap.MarkReadyForSession(context.Background(), msg.OwnerUserID, authKeyID, newSessionID) + if err != nil || ready != 1 { + t.Fatalf("same auth-key reconnect ready=%d err=%v, want 1/nil", ready, err) + } + if claimed := r.publishReadyBootstrapUpdates(context.Background(), 1, time.Second, zaptest.NewLogger(t)); claimed != 1 { + t.Fatalf("published after same-auth reconnect = %d, want 1", claimed) + } + if len(updates.events) != 1 || updates.events[0].Message.ID != msg.ID { + t.Fatalf("events after reconnect = %+v", updates.events) + } +} + +func TestUpdatesGetStatePublishesSignUpBootstrapAfterRPCResult(t *testing.T) { bootstrap := memory.NewBootstrapUpdateJobStore() updates := &captureUpdates{state: domain.UpdateState{Pts: 0, Date: 1700000000}} msg := domain.Message{ @@ -86,9 +123,12 @@ func TestUpdatesGetStatePublishesBootstrapAfterRPCResult(t *testing.T) { authKeyID := [8]byte{9, 8, 7} sessionID := int64(5723482677041206318) ctx := postresponse.WithCallbacks( - WithUserID( - WithSessionID(WithAuthKeyID(context.Background(), authKeyID), sessionID), - msg.OwnerUserID, + WithClientInfo( + WithUserID( + WithSessionID(WithAuthKeyID(context.Background(), authKeyID), sessionID), + msg.OwnerUserID, + ), + ClientInfo{Type: ClientTypeTDesktop}, ), ) r.enqueueLoginMessageBootstrap(ctx, msg) @@ -180,7 +220,7 @@ func TestLogOutClearsSessionAndUpdateState(t *testing.T) { } } -func TestSignInDifferentUserClearsAuthKeyUpdateState(t *testing.T) { +func TestSignInDifferentUserDoesNotClearFreshlyBoundUpdateState(t *testing.T) { var authKeyID [8]byte authKeyID[0] = 6 auth := &captureAuthService{signInUser: domain.User{ID: 1000000002, FirstName: "Two"}} @@ -197,8 +237,8 @@ func TestSignInDifferentUserClearsAuthKeyUpdateState(t *testing.T) { if err != nil { t.Fatalf("auth.signIn: %v", err) } - if updates.clearedAuthKeyID != authKeyID || !updates.cleared { - t.Fatalf("cleared auth key = %x cleared=%v, want %x", updates.clearedAuthKeyID, updates.cleared, authKeyID) + if updates.cleared { + t.Fatalf("router cleared auth key %x after Bind; this would delete the new user's retained-floor baseline", updates.clearedAuthKeyID) } gotSession := sessions.snapshot() if gotSession.userID != 1000000002 { @@ -213,7 +253,8 @@ func TestUpdatesGetStateMarksSessionReadyForPush(t *testing.T) { Updates: &captureUpdates{state: domain.UpdateState{Pts: 3, Date: 1700000000, Seq: 2}}, }, zaptest.NewLogger(t), clock.System) - got, err := r.onUpdatesGetState(WithSessionID(context.Background(), 77)) + ctx := WithClientInfo(WithSessionID(context.Background(), 77), ClientInfo{Type: ClientTypeTDesktop}) + got, err := r.onUpdatesGetState(ctx) if err != nil { t.Fatalf("updates.getState: %v", err) } @@ -243,7 +284,8 @@ func TestUpdatesGetStateReturnsAccountCurrentState(t *testing.T) { Updates: updates, }, zaptest.NewLogger(t), clock.System) - got, err := r.onUpdatesGetState(WithSessionID(context.Background(), 77)) + ctx := WithClientInfo(WithSessionID(context.Background(), 77), ClientInfo{Type: ClientTypeTDesktop}) + got, err := r.onUpdatesGetState(ctx) if err != nil { t.Fatalf("updates.getState: %v", err) } @@ -261,6 +303,43 @@ func TestUpdatesGetStateReturnsAccountCurrentState(t *testing.T) { } } +func TestUpdatesGetStateUnknownClientDoesNotAdvanceObservedBaseline(t *testing.T) { + sessions := &captureSessions{} + current := domain.UpdateState{Pts: 9, Date: 1700000009} + updates := &captureUpdates{ + state: domain.UpdateState{Pts: 3, Date: 1700000003}, + currentState: ¤t, + } + r := New(Config{}, Deps{Sessions: sessions, Updates: updates}, zaptest.NewLogger(t), clock.System) + + got, err := r.onUpdatesGetState(WithSessionID(context.Background(), 78)) + if err != nil { + t.Fatalf("updates.getState unknown client: %v", err) + } + if got.Pts != current.Pts { + t.Fatalf("state pts = %d, want current %d", got.Pts, current.Pts) + } + if updates.acknowledged { + t.Fatal("unknown client advanced the observed getState baseline") + } + if !sessions.snapshot().receives { + t.Fatal("unknown client was not enabled for subsequent updates") + } +} + +func TestUpdatesGetStateDrKLOEstablishesObservedBaseline(t *testing.T) { + updates := &captureUpdates{currentState: &domain.UpdateState{Pts: 6, Date: 1700000006}} + r := New(Config{}, Deps{Sessions: &captureSessions{}, Updates: updates}, zaptest.NewLogger(t), clock.System) + ctx := WithClientInfo(context.Background(), ClientInfo{Type: ClientTypeAndroid, AppVersion: "12.8.1"}) + + if _, err := r.onUpdatesGetState(ctx); err != nil { + t.Fatalf("updates.getState DrKLO: %v", err) + } + if !updates.acknowledged { + t.Fatal("DrKLO getState did not establish its explicit current-snapshot baseline") + } +} + func TestUpdatesDifferenceIncludesLoginMessageAndOfficialUser(t *testing.T) { msg := domain.Message{ ID: 88, @@ -575,6 +654,28 @@ func TestUpdatesGetDifferenceChannelNudgeIncludesFullChat(t *testing.T) { } } +func TestUpdatesGetDifferenceEmptyRetentionCheckpointStaysDifferenceSlice(t *testing.T) { + updates := &captureUpdates{difference: &domain.UpdateDifference{ + State: domain.UpdateState{Pts: 42, Date: 1700000442}, + Partial: true, + }} + r := New(Config{}, Deps{Updates: updates}, zaptest.NewLogger(t), fixedClock{now: time.Unix(1700000443, 0)}) + authKeyID := [8]byte{42} + ctx := WithUserID(WithAuthKeyID(context.Background(), authKeyID), 1000000042) + + diff, err := r.onUpdatesGetDifference(ctx, &tg.UpdatesGetDifferenceRequest{Pts: 1, Date: 1700000400}) + if err != nil { + t.Fatalf("updates.getDifference: %v", err) + } + slice, ok := diff.(*tg.UpdatesDifferenceSlice) + if !ok { + t.Fatalf("difference = %T, want *tg.UpdatesDifferenceSlice (never Empty/TooLong)", diff) + } + if slice.IntermediateState.Pts != 42 || slice.IntermediateState.Date != 1700000442 || len(slice.NewMessages) != 0 || len(slice.OtherUpdates) != 0 { + t.Fatalf("checkpoint slice = %+v, want empty payload at pts/date 42/1700000442", slice) + } +} + func TestUpdatesDifferenceIncludesSettingsUpdates(t *testing.T) { peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002} got, ok := tgUpdatesDifference(0, domain.UpdateDifference{ diff --git a/internal/store/album_group.go b/internal/store/album_group.go new file mode 100644 index 00000000..48347918 --- /dev/null +++ b/internal/store/album_group.go @@ -0,0 +1,18 @@ +package store + +import ( + "context" + + "telesrv/internal/domain" +) + +// AlbumGroupStore 持久化 sendMultiMedia 的预发送分组预留。 +// +// ReserveAlbumGroup 必须原子满足： +// - 请求内没有旧绑定时，全部 random_id 绑定 ProposedGroupedID； +// - 命中唯一旧 grouped_id 时，全部缺失项收敛到该旧值； +// - 命中多个旧 grouped_id 时返回 domain.ErrMessageRandomIDDuplicate，且不写入； +// - 并发、多实例的重叠请求等价于某个串行顺序。 +type AlbumGroupStore interface { + ReserveAlbumGroup(ctx context.Context, req domain.AlbumGroupReservationRequest) (groupedID int64, err error) +} diff --git a/internal/store/bootstrap_update_job.go b/internal/store/bootstrap_update_job.go index f6693a5a..65a90236 100644 --- a/internal/store/bootstrap_update_job.go +++ b/internal/store/bootstrap_update_job.go @@ -9,6 +9,8 @@ import ( type BootstrapUpdateJobStore interface { EnqueueLoginMessage(ctx context.Context, job domain.BootstrapUpdateJob) (domain.BootstrapUpdateJob, error) + // MarkReadyForSession allows a new session on the same auth key to take over + // a pending signup baseline fence; a different auth key can never release it. MarkReadyForSession(ctx context.Context, userID int64, authKeyID [8]byte, sessionID int64) (int, error) ClaimReady(ctx context.Context, limit int, leaseTimeout time.Duration) ([]domain.BootstrapUpdateJob, error) MarkPublished(ctx context.Context, id int64) error diff --git a/internal/store/channel.go b/internal/store/channel.go index b758f7d3..deaf06f2 100644 --- a/internal/store/channel.go +++ b/internal/store/channel.go @@ -2,6 +2,7 @@ package store import ( "context" + "time" "telesrv/internal/domain" ) @@ -164,6 +165,12 @@ type ChannelStore interface { GeneralForumTopic(ctx context.Context, viewerUserID, channelID int64) (domain.ChannelForumTopic, error) ListMessageReadParticipants(ctx context.Context, req domain.ChannelReadParticipantsRequest) (domain.ChannelReadParticipantsResult, error) ListChannelDifference(ctx context.Context, req domain.ChannelDifferenceRequest) (domain.ChannelDifference, error) + // PruneChannelUpdateEvents atomically removes at most limit complete event rows through throughPts + // and advances the channel retained floor only to the last contiguous row actually removed. + PruneChannelUpdateEvents(ctx context.Context, channelID int64, throughPts, limit int) (domain.ChannelUpdateRetentionResult, error) + // DeleteExpiredChannelUpdateEvents selects expired channel-log heads using an indexed, bounded seek + // and delegates each channel to the same atomic floor+delete primitive. It never uses SQL OFFSET. + DeleteExpiredChannelUpdateEvents(ctx context.Context, olderThan time.Duration, limit int) (int, error) ListActiveChannelIDsForUser(ctx context.Context, userID, afterChannelID int64, limit int) ([]int64, error) ListDirtyActiveChannelsForUser(ctx context.Context, userID int64, sinceDate int, afterChannelID int64, limit int) ([]domain.DirtyChannel, error) ListActiveChannelMemberIDs(ctx context.Context, viewerUserID, channelID int64, limit int) ([]int64, error) @@ -171,6 +178,10 @@ type ChannelStore interface { ListChannelInviteAdminMemberIDs(ctx context.Context, channelID int64, limit int) ([]int64, error) FilterActiveChannelMemberIDs(ctx context.Context, channelID int64, userIDs []int64) ([]int64, error) MaxChannelPts(ctx context.Context, channelID int64) (int, error) + // MaxChannelPtsBatch returns existing channel watermarks with one bounded store round trip. + // Missing/deleted ids are omitted so a stale process-local membership key cannot poison the + // entire fan-out recovery sweep. + MaxChannelPtsBatch(ctx context.Context, channelIDs []int64) (map[int64]int, error) // SetActiveCall 写入/清除（callID=0）channel 行上的活跃群通话关联 //（channel.call_active/call_not_empty flag 与 channelFull.call 的数据源）。 SetActiveCall(ctx context.Context, channelID, callID, callAccessHash int64, notEmpty bool) (domain.Channel, error) diff --git a/internal/store/code.go b/internal/store/code.go index 6bf75c3c..72b34eba 100644 --- a/internal/store/code.go +++ b/internal/store/code.go @@ -2,21 +2,47 @@ package store import ( "context" + "crypto/rand" + "encoding/hex" + "fmt" "time" ) -const PhoneCodePurposeChangePhone = "change_phone" +const ( + PhoneCodePurposeChangePhone = "change_phone" + PhoneCodeChannelPhone = "phone" + PhoneCodeChannelEmailLogin = "email_login" + PhoneCodeChannelEmailSetupRequired = "email_setup_required" +) + +// PhoneCodeVersionCurrent is the only version accepted by the atomic login +// state machine. Version zero is the pre-state-machine shape and deliberately +// fails closed instead of being normalized on read. +const PhoneCodeVersionCurrent = 1 // PhoneCode 是一条验证码记录（与某次 sendCode 的 phone_code_hash 或邮箱验证键关联）。 // Purpose/UserID/AuthKeyID/SessionID 为已登录敏感操作提供作用域；登录验证码保持零值。 type PhoneCode struct { + Version int + // Revision is an opaque store-managed CAS token. Callers must pass it back + // through PhoneCodeSnapshot and must never synthesize or persist it outside + // CodeStore. + Revision string + // IssuedUserID is encoded as a JSON string so Redis Lua can round-trip the + // full int64 range without cjson's IEEE-754 number precision loss. + IssuedUserID int64 `json:",string"` + SignUpVerified bool Phone string Code string Channel string Purpose string - UserID int64 - AuthKeyID [8]byte - SessionID int64 + // UserID is also encoded as a string because scoped verification mutates the + // record in Redis Lua and must not round an int64 owner through cjson. + UserID int64 `json:",string"` + AuthKeyID [8]byte + // SessionID is audit metadata but still crosses Redis Lua on wrong attempts; + // encode it as a string to preserve MTProto's full signed 64-bit value. + SessionID int64 `json:",string"` Email string PendingEmail string Attempts int @@ -26,6 +52,35 @@ type PhoneCode struct { LoginEmailHash string } +type PhoneCodeSnapshot struct { + Record PhoneCode + Revision string +} + +func NewPhoneCodeRevisionToken() (string, error) { + var raw [16]byte + if _, err := rand.Read(raw[:]); err != nil { + return "", fmt.Errorf("generate phone code revision: %w", err) + } + return hex.EncodeToString(raw[:]), nil +} + +// LoginCodeVerifyStatus separates an expired/consumed hash from a live hash +// whose scope or code did not match. RPC maps these to PHONE_CODE_EXPIRED and +// PHONE_CODE_INVALID respectively. +type LoginCodeVerifyStatus uint8 + +const ( + LoginCodeVerifyMissing LoginCodeVerifyStatus = iota + LoginCodeVerifyInvalid + LoginCodeVerifyAccepted +) + +type LoginCodeVerifyResult struct { + Status LoginCodeVerifyStatus + Record PhoneCode +} + // PhoneCodeScope 标识已登录敏感操作的一次性验证码作用域。SessionID 故意不在 // 作用域内：同一 perm auth key 等待验证码期间允许重建 MTProto session。 // 登录/注册验证码没有 Purpose/UserID/AuthKeyID，保持非 scoped 行为。 @@ -56,9 +111,36 @@ type CodeStore interface { // 活跃验证码；普通登录码仍按 hash 独立保存。 Set(ctx context.Context, phoneCodeHash string, code PhoneCode, ttl time.Duration) error Get(ctx context.Context, phoneCodeHash string) (PhoneCode, bool, error) - Update(ctx context.Context, phoneCodeHash string, code PhoneCode) error Del(ctx context.Context, phoneCodeHash string) error // ConsumeScoped 仅当 hash 仍是 scope 的当前活跃 hash 时原子读取并删除； // 并发调用至多一个返回 found=true。 ConsumeScoped(ctx context.Context, phoneCodeHash string, scope PhoneCodeScope) (PhoneCode, bool, error) + // VerifyScoped atomically verifies a current-version code only while hash is + // still the active value for scope. A wrong code increments Attempts and, at + // the threshold, removes both the code and scope index. A correct code + // consumes both keys, so concurrent callers can observe Accepted at most once. + VerifyScoped(ctx context.Context, phoneCodeHash string, scope PhoneCodeScope, code string, defaultMaxAttempts int) (LoginCodeVerifyResult, error) + // VerifyLogin atomically validates one current-version, unscoped login code. + // A correct code is consumed unless keepForSignUp is true, in which case the + // same TTL is retained and SignUpVerified is set. Wrong-code attempts are + // incremented in the same linearization point and delete the record at the + // configured threshold. + VerifyLogin(ctx context.Context, phoneCodeHash, phone, code string, keepForSignUp bool, defaultMaxAttempts int) (LoginCodeVerifyResult, error) + // ConsumeSignUpVerified atomically consumes a marker created by VerifyLogin. + // Concurrent sign-up calls can return found=true at most once. + ConsumeSignUpVerified(ctx context.Context, phoneCodeHash, phone string) (PhoneCode, bool, error) + // TakeLoginCode atomically removes a current-version, unscoped login record + // after matching its phone. Cancel/resend use the returned record to decide + // what successor to issue; concurrent Verify/Take calls have one winner. + TakeLoginCode(ctx context.Context, phoneCodeHash, phone string) (PhoneCode, bool, error) + // InvalidateLoginCode is the server-side cleanup primitive for owner drift + // or a failed post-verification workflow. Unlike TakeLoginCode it may delete + // a SignUpVerified marker; user-driven cancel/resend must not call it. + InvalidateLoginCode(ctx context.Context, phoneCodeHash, phone string) (bool, error) + // GetSnapshot and CompareAnd* provide optimistic concurrency for unscoped + // fixed-key verification flows (for example login-email setup/change). + // CompareAndUpdate preserves the current TTL and rotates the opaque revision. + GetSnapshot(ctx context.Context, phoneCodeHash string) (PhoneCodeSnapshot, bool, error) + CompareAndUpdate(ctx context.Context, phoneCodeHash, expectedRevision string, next PhoneCode) (bool, error) + CompareAndDelete(ctx context.Context, phoneCodeHash, expectedRevision string) (bool, error) } diff --git a/internal/store/dispatch_outbox.go b/internal/store/dispatch_outbox.go index 6f71cbdf..f83fab31 100644 --- a/internal/store/dispatch_outbox.go +++ b/internal/store/dispatch_outbox.go @@ -2,11 +2,20 @@ package store import ( "context" + "errors" "time" "telesrv/internal/domain" ) +// ErrDispatchLeaseLost means a completion belongs to an older claim attempt. +// Callers must not overwrite/delete the row now owned by a newer worker. +var ErrDispatchLeaseLost = errors.New("dispatch outbox lease lost") + +// DispatchOutboxLogicalShards 是稳定的 user→lane 哈希空间。运行时 worker 数 +// 只能改变 shard 的归属，不能改变这个值；PG 表达式索引也固定使用 256。 +const DispatchOutboxLogicalShards = 256 + // DispatchOutboxItem 是待投递给在线 session 的 update 任务。 type DispatchOutboxItem struct { ID int64 @@ -21,7 +30,7 @@ type DispatchOutboxItem struct { // DispatchOutboxStore 持久化 transactional outbox。 type DispatchOutboxStore interface { ClaimPending(ctx context.Context, limit int) ([]DispatchOutboxItem, error) - MarkDelivered(ctx context.Context, targetUserID, id int64) error - MarkFailed(ctx context.Context, targetUserID, id int64, lastError string) error + MarkDelivered(ctx context.Context, item DispatchOutboxItem) error + MarkFailed(ctx context.Context, item DispatchOutboxItem, lastError string) error DeleteFailed(ctx context.Context, olderThan time.Duration, limit int) (int, error) } diff --git a/internal/store/login_code_delivery.go b/internal/store/login_code_delivery.go new file mode 100644 index 00000000..8e08bfaa --- /dev/null +++ b/internal/store/login_code_delivery.go @@ -0,0 +1,67 @@ +package store + +import ( + "context" + "crypto/hmac" + "crypto/sha256" + "crypto/subtle" + "fmt" + + "telesrv/internal/domain" +) + +const ( + maxPhoneCodeHashBytes = 512 + maxLoginCodeBytes = 64 +) + +// LoginCodeDeliveryStore atomically creates the recipient message box, dialog, +// account update event and online-dispatch task for a 777000 login code. +type LoginCodeDeliveryStore interface { + DeliverLoginCodeMessage(ctx context.Context, req domain.LoginCodeDeliveryRequest) (domain.LoginCodeDeliveryResult, error) +} + +// LoginCodeDeliveryKey is the only phone-code-hash representation permitted at +// rest. The raw phone_code_hash stays at the auth/store call boundary. +func LoginCodeDeliveryKey(phoneCodeHash string) ([sha256.Size]byte, error) { + if phoneCodeHash == "" || len(phoneCodeHash) > maxPhoneCodeHashBytes { + return [sha256.Size]byte{}, domain.ErrLoginCodeDeliveryInvalid + } + return sha256.Sum256([]byte(phoneCodeHash)), nil +} + +// LoginCodeFingerprint binds a delivery receipt to its immutable secret code. +// It is keyed by the high-entropy raw phone_code_hash, which is never stored; +// unlike a bare digest of a short login code this cannot be brute-forced from +// the compact receipt alone after the message itself has been deleted. +func LoginCodeFingerprint(phoneCodeHash, code string) ([sha256.Size]byte, error) { + if phoneCodeHash == "" || len(phoneCodeHash) > maxPhoneCodeHashBytes || code == "" || len(code) > maxLoginCodeBytes { + return [sha256.Size]byte{}, domain.ErrLoginCodeDeliveryInvalid + } + mac := hmac.New(sha256.New, []byte(phoneCodeHash)) + _, _ = mac.Write([]byte(code)) + var fingerprint [sha256.Size]byte + copy(fingerprint[:], mac.Sum(nil)) + return fingerprint, nil +} + +func SameLoginCodeFingerprint(stored []byte, expected [sha256.Size]byte) bool { + return len(stored) == sha256.Size && subtle.ConstantTimeCompare(stored, expected[:]) == 1 +} + +// RestoreLoginCodeDeliveryMessage reconstructs the immutable first result from +// a compact receipt. The secret code is not duplicated in the receipt: exact +// replay has already proven the supplied code fingerprint matches. +func RestoreLoginCodeDeliveryMessage(userID int64, code string, date int, privateMessageID int64, messageBoxID, pts int) (domain.Message, error) { + if privateMessageID <= 0 || messageBoxID <= 0 || messageBoxID > domain.MaxMessageBoxID || pts <= 0 { + return domain.Message{}, fmt.Errorf("restore login code delivery: %w: uid=%d box=%d pts=%d", domain.ErrLoginCodeDeliveryInvalid, privateMessageID, messageBoxID, pts) + } + msg, err := domain.OfficialLoginCodeMessage(userID, code, date) + if err != nil { + return domain.Message{}, err + } + msg.ID = messageBoxID + msg.UID = privateMessageID + msg.Pts = pts + return msg, nil +} diff --git a/internal/store/login_code_delivery_test.go b/internal/store/login_code_delivery_test.go new file mode 100644 index 00000000..6396c553 --- /dev/null +++ b/internal/store/login_code_delivery_test.go @@ -0,0 +1,67 @@ +package store + +import ( + "crypto/sha256" + "errors" + "reflect" + "testing" + + "telesrv/internal/domain" +) + +func TestLoginCodeDeliveryKeyAndFingerprint(t *testing.T) { + const phoneCodeHash = "opaque-high-entropy-phone-code-hash" + key, err := LoginCodeDeliveryKey(phoneCodeHash) + if err != nil { + t.Fatalf("LoginCodeDeliveryKey: %v", err) + } + if want := sha256.Sum256([]byte(phoneCodeHash)); key != want { + t.Fatalf("delivery key = %x, want SHA-256 %x", key, want) + } + + fingerprint, err := LoginCodeFingerprint(phoneCodeHash, "12345") + if err != nil { + t.Fatalf("LoginCodeFingerprint: %v", err) + } + if !SameLoginCodeFingerprint(fingerprint[:], fingerprint) { + t.Fatal("fingerprint does not compare equal to itself") + } + otherHash, err := LoginCodeFingerprint("another-phone-code-hash", "12345") + if err != nil { + t.Fatalf("other hash fingerprint: %v", err) + } + otherCode, err := LoginCodeFingerprint(phoneCodeHash, "54321") + if err != nil { + t.Fatalf("other code fingerprint: %v", err) + } + if fingerprint == otherHash || fingerprint == otherCode { + t.Fatal("fingerprint must bind both the raw phone_code_hash and code") + } + if SameLoginCodeFingerprint(fingerprint[:31], fingerprint) { + t.Fatal("truncated fingerprint compared equal") + } + if _, err := LoginCodeDeliveryKey(string(make([]byte, maxPhoneCodeHashBytes+1))); !errors.Is(err, domain.ErrLoginCodeDeliveryInvalid) { + t.Fatalf("oversized phone_code_hash err = %v, want ErrLoginCodeDeliveryInvalid", err) + } + if _, err := LoginCodeFingerprint(phoneCodeHash, string(make([]byte, maxLoginCodeBytes+1))); !errors.Is(err, domain.ErrLoginCodeDeliveryInvalid) { + t.Fatalf("oversized code err = %v, want ErrLoginCodeDeliveryInvalid", err) + } +} + +func TestRestoreLoginCodeDeliveryMessage(t *testing.T) { + got, err := RestoreLoginCodeDeliveryMessage(1000000001, "12345", 1700000000, 91, 7, 12) + if err != nil { + t.Fatalf("RestoreLoginCodeDeliveryMessage: %v", err) + } + want, err := domain.OfficialLoginCodeMessage(1000000001, "12345", 1700000000) + if err != nil { + t.Fatalf("OfficialLoginCodeMessage: %v", err) + } + want.UID, want.ID, want.Pts = 91, 7, 12 + if !reflect.DeepEqual(got, want) { + t.Fatalf("restored message = %+v, want %+v", got, want) + } + if _, err := RestoreLoginCodeDeliveryMessage(1000000001, "12345", 1700000000, 0, 7, 12); !errors.Is(err, domain.ErrLoginCodeDeliveryInvalid) { + t.Fatalf("invalid uid err = %v, want ErrLoginCodeDeliveryInvalid", err) + } +} diff --git a/internal/store/media.go b/internal/store/media.go index f79b37be..0db7ceae 100644 --- a/internal/store/media.go +++ b/internal/store/media.go @@ -18,6 +18,11 @@ type MediaStore interface { LoadFileParts(ctx context.Context, ownerUserID, fileID int64) ([]domain.UploadPart, error) DeleteFileParts(ctx context.Context, ownerUserID, fileID int64) ([]string, error) DeleteExpiredUploadParts(ctx context.Context, before time.Time, limit int) ([]string, error) + // Uploaded media receipts survive transient part cleanup so messages.sendMedia can replay an + // InputMediaUploadedPhoto/Document after a lost response. Put returns the durable winner; when + // created=false another concurrent materializer won the (owner,file_id) key. + GetUploadedMediaReceipt(ctx context.Context, ownerUserID, fileID int64) (domain.UploadedMediaReceipt, bool, error) + PutUploadedMediaReceipt(ctx context.Context, receipt domain.UploadedMediaReceipt) (stored domain.UploadedMediaReceipt, created bool, err error) // blob 索引。 PutFileBlob(ctx context.Context, blob domain.FileBlob) error diff --git a/internal/store/memory/album_group.go b/internal/store/memory/album_group.go new file mode 100644 index 00000000..601be2e2 --- /dev/null +++ b/internal/store/memory/album_group.go @@ -0,0 +1,68 @@ +package memory + +import ( + "bytes" + "context" + + "telesrv/internal/domain" +) + +type albumGroupKey struct { + senderUserID int64 + peerType domain.PeerType + peerID int64 + randomID int64 +} + +type albumGroupRecord struct { + groupedID int64 + intentHash [32]byte +} + +// ReserveAlbumGroup 在 MessageStore 的同一把互斥锁下完成读旧组、选胜者与补齐绑定， +// 因而并发重叠批次不会产生拆组。 +func (s *MessageStore) ReserveAlbumGroup(_ context.Context, req domain.AlbumGroupReservationRequest) (int64, error) { + if err := req.Validate(); err != nil { + return 0, err + } + + s.mu.Lock() + defer s.mu.Unlock() + if s.albumGroups == nil { + s.albumGroups = make(map[albumGroupKey]albumGroupRecord) + } + + groupedID := int64(0) + type pendingBinding struct { + key albumGroupKey + intentHash [32]byte + } + bindings := make([]pendingBinding, 0, len(req.Items)) + for _, item := range req.Items { + key := albumGroupKey{ + senderUserID: req.SenderUserID, + peerType: req.Peer.Type, + peerID: req.Peer.ID, + randomID: item.RandomID, + } + var intentHash [32]byte + copy(intentHash[:], item.IntentHash) + bindings = append(bindings, pendingBinding{key: key, intentHash: intentHash}) + if existing, exists := s.albumGroups[key]; exists { + if !bytes.Equal(existing.intentHash[:], item.IntentHash) { + return 0, domain.ErrMessageRandomIDDuplicate + } + if groupedID != 0 && groupedID != existing.groupedID { + return 0, domain.ErrMessageRandomIDDuplicate + } + groupedID = existing.groupedID + } + } + if groupedID == 0 { + groupedID = req.ProposedGroupedID + } + for _, binding := range bindings { + s.albumGroups[binding.key] = albumGroupRecord{groupedID: groupedID, intentHash: binding.intentHash} + } + return groupedID, nil +} diff --git a/internal/store/memory/album_group_test.go b/internal/store/memory/album_group_test.go new file mode 100644 index 00000000..9df60abf --- /dev/null +++ b/internal/store/memory/album_group_test.go @@ -0,0 +1,131 @@ +package memory + +import ( + "context" + "crypto/sha256" + "errors" + "sync" + "testing" + + "telesrv/internal/domain" +) + +func albumIntent(label string) []byte { + sum := sha256.Sum256([]byte(label)) + return sum[:] +} + +func albumReq(sender int64, peer domain.Peer, groupedID int64, items ...domain.AlbumGroupReservationItem) domain.AlbumGroupReservationRequest { + return domain.AlbumGroupReservationRequest{ + SenderUserID: sender, + Peer: peer, + Items: items, + ProposedGroupedID: groupedID, + } +} + +func albumItem(randomID int64, label string) domain.AlbumGroupReservationItem { + return domain.AlbumGroupReservationItem{RandomID: randomID, IntentHash: albumIntent(label)} +} + +func TestAlbumGroupReservationFullThenSubsetAndIntentConflict(t *testing.T) { + ctx := context.Background() + messages := NewMessageStore() + peer := domain.Peer{Type: domain.PeerTypeUser, ID: 2002} + full := []domain.AlbumGroupReservationItem{albumItem(1, "one"), albumItem(2, "two"), albumItem(3, "three")} + + groupedID, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 101, full...)) + if err != nil || groupedID != 101 { + t.Fatalf("reserve full = %d err=%v, want 101", groupedID, err) + } + replayed, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 202, full[1:]...)) + if err != nil || replayed != groupedID { + t.Fatalf("reserve subset = %d err=%v, want original %d", replayed, err, groupedID) + } + for _, item := range full { + got, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 303, item)) + if err != nil || got != groupedID { + t.Fatalf("single random_id %d = %d err=%v, want %d", item.RandomID, got, err, groupedID) + } + } + changed := albumItem(2, "changed payload") + if _, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 404, changed)); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("changed intent err=%v, want ErrMessageRandomIDDuplicate", err) + } +} + +func TestAlbumGroupReservationConcurrentOverlapConverges(t *testing.T) { + ctx := context.Background() + messages := NewMessageStore() + peer := domain.Peer{Type: domain.PeerTypeChannel, ID: 9001} + requests := []domain.AlbumGroupReservationRequest{ + albumReq(1001, peer, 111, albumItem(11, "one"), albumItem(12, "shared")), + albumReq(1001, peer, 222, albumItem(12, "shared"), albumItem(13, "three")), + } + results := make([]int64, 2) + errs := make([]error, 2) + start := make(chan struct{}) + var wg sync.WaitGroup + for i := range requests { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + results[i], errs[i] = messages.ReserveAlbumGroup(ctx, requests[i]) + }(i) + } + close(start) + wg.Wait() + if errs[0] != nil || errs[1] != nil || results[0] == 0 || results[0] != results[1] { + t.Fatalf("concurrent results=%v errs=%v, want same non-zero group", results, errs) + } + for _, item := range []domain.AlbumGroupReservationItem{albumItem(11, "one"), albumItem(12, "shared"), albumItem(13, "three")} { + got, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 333, item)) + if err != nil || got != results[0] { + t.Fatalf("converged random_id %d = %d err=%v, want %d", item.RandomID, got, err, results[0]) + } + } +} + +func TestAlbumGroupReservationRejectsMixedOldGroupsAtomically(t *testing.T) { + ctx := context.Background() + messages := NewMessageStore() + peer := domain.Peer{Type: domain.PeerTypeUser, ID: 2002} + one := albumItem(21, "one") + two := albumItem(22, "two") + three := albumItem(23, "three") + if _, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 121, one)); err != nil { + t.Fatal(err) + } + if _, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 122, two)); err != nil { + t.Fatal(err) + } + if _, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 123, one, two, three)); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("mixed old groups err=%v, want ErrMessageRandomIDDuplicate", err) + } + // 失败批次不能把尚未存在的 random_id 23 偷绑到任一旧组。 + got, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, peer, 124, three)) + if err != nil || got != 124 { + t.Fatalf("post-conflict unbound item = %d err=%v, want fresh 124", got, err) + } +} + +func TestAlbumGroupReservationScopeIncludesPeer(t *testing.T) { + ctx := context.Background() + messages := NewMessageStore() + item := albumItem(31, "same intent") + tests := []struct { + peer domain.Peer + group int64 + }{ + {peer: domain.Peer{Type: domain.PeerTypeUser, ID: 2002}, group: 131}, + {peer: domain.Peer{Type: domain.PeerTypeUser, ID: 2003}, group: 132}, + {peer: domain.Peer{Type: domain.PeerTypeChannel, ID: 2002}, group: 133}, + } + for _, tc := range tests { + got, err := messages.ReserveAlbumGroup(ctx, albumReq(1001, tc.peer, tc.group, item)) + if err != nil || got != tc.group { + t.Fatalf("peer %+v = %d err=%v, want isolated %d", tc.peer, got, err, tc.group) + } + } +} diff --git a/internal/store/memory/auth.go b/internal/store/memory/auth.go index 895e2c48..e0fc4ec6 100644 --- a/internal/store/memory/auth.go +++ b/internal/store/memory/auth.go @@ -276,6 +276,11 @@ func NewCodeStore() *CodeStore { } func (s *CodeStore) Set(_ context.Context, hash string, code store.PhoneCode, ttl time.Duration) error { + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return err + } + code.Revision = revision s.mu.Lock() scope := code.Scope() if scope.Valid() { @@ -303,6 +308,11 @@ func (s *CodeStore) Get(_ context.Context, hash string) (store.PhoneCode, bool, } func (s *CodeStore) Update(_ context.Context, hash string, code store.PhoneCode) error { + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return err + } + code.Revision = revision s.mu.Lock() defer s.mu.Unlock() e, ok := s.m[hash] @@ -343,7 +353,13 @@ func (s *CodeStore) ConsumeScoped(_ context.Context, hash string, scope store.Ph } return store.PhoneCode{}, false, nil } - if e.code.Scope() != scope { + if e.code.Version != store.PhoneCodeVersionCurrent || e.code.Scope() != scope { + delete(s.m, hash) + delete(s.scopes, scope) + actualScope := e.code.Scope() + if actualScope.Valid() && s.scopes[actualScope] == hash { + delete(s.scopes, actualScope) + } return store.PhoneCode{}, false, nil } s.deleteCodeLocked(hash, e.code) diff --git a/internal/store/memory/bootstrap_update_job.go b/internal/store/memory/bootstrap_update_job.go index 7132b7df..d5cddd63 100644 --- a/internal/store/memory/bootstrap_update_job.go +++ b/internal/store/memory/bootstrap_update_job.go @@ -67,10 +67,15 @@ func (s *BootstrapUpdateJobStore) MarkReadyForSession(_ context.Context, userID s.mu.Lock() defer s.mu.Unlock() for id, job := range s.jobs { - if job.UserID != userID || job.AuthKeyID != authKeyID || job.SessionID != sessionID || job.Status != domain.BootstrapUpdateJobPending { + if job.UserID != userID || job.AuthKeyID != authKeyID || job.Status != domain.BootstrapUpdateJobPending { continue } job.Status = domain.BootstrapUpdateJobReady + // A reconnect on the same physical/business auth key is the same + // verified device. Transfer the pending baseline fence to the session + // that actually completed getState/getDifference instead of orphaning + // the job on the disconnected signup session. + job.SessionID = sessionID job.ReadyAt = now job.UpdatedAt = now s.jobs[id] = job diff --git a/internal/store/memory/channel_helpers.go b/internal/store/memory/channel_helpers.go index c8d9fa5b..5e99d59f 100644 --- a/internal/store/memory/channel_helpers.go +++ b/internal/store/memory/channel_helpers.go @@ -319,14 +319,8 @@ func (s *ChannelStore) ListDirtyActiveChannelsForUser(_ context.Context, userID if !ok || member.Status != domain.ChannelMemberActive { continue } - dirty := false - for _, event := range s.events[channelID] { - if event.Date > sinceDate { - dirty = true - break - } - } - if dirty { + checkpoint := s.channelUpdateCheckpointLocked(channelID, channel) + if checkpoint.LatestEventDate > sinceDate { out = append(out, domain.DirtyChannel{ChannelID: channelID, Pts: channel.Pts}) } } diff --git a/internal/store/memory/channel_message_delete.go b/internal/store/memory/channel_message_delete.go index 03355c25..06831b16 100644 --- a/internal/store/memory/channel_message_delete.go +++ b/internal/store/memory/channel_message_delete.go @@ -240,7 +240,16 @@ func (s *ChannelStore) deleteChannelMessagesLocked(channel domain.Channel, membe MessageIDs: append([]int(nil), deleted...), SenderUserID: actorUserID, } - s.events[channel.ID] = append(s.events[channel.ID], event) + s.appendChannelEventLocked(event) + for _, id := range deleted { + msg, ok := s.findMessageLocked(channel.ID, id) + if !ok || msg.RandomID == 0 || msg.SenderUserID == 0 { + continue + } + key := channelMessageReplayKey{channelID: channel.ID, messageID: msg.ID} + cloned := cloneChannelEvent(event) + s.deleteReceipts[key] = &cloned + } return deleted, event, channel, nil } diff --git a/internal/store/memory/channel_message_edit.go b/internal/store/memory/channel_message_edit.go index db3862d3..3edfe52e 100644 --- a/internal/store/memory/channel_message_edit.go +++ b/internal/store/memory/channel_message_edit.go @@ -58,7 +58,7 @@ func (s *ChannelStore) EditChannelMessage(_ context.Context, req domain.EditChan Message: cloneChannelMessage(msg), SenderUserID: req.UserID, } - s.events[req.ChannelID] = append(s.events[req.ChannelID], event) + s.appendChannelEventLocked(event) return domain.EditChannelMessageResult{ Channel: channel, Message: cloneChannelMessage(msg), @@ -108,7 +108,7 @@ func (s *ChannelStore) EditChannelMessage(_ context.Context, req domain.EditChan Message: cloneChannelMessage(msg), SenderUserID: req.UserID, } - s.events[req.ChannelID] = append(s.events[req.ChannelID], event) + s.appendChannelEventLocked(event) s.appendChannelAdminLogLocked(domain.ChannelAdminLogEvent{ ChannelID: req.ChannelID, UserID: req.UserID, @@ -150,7 +150,7 @@ func (s *ChannelStore) EditChannelMessage(_ context.Context, req domain.EditChan SenderUserID: req.UserID, } s.messages[req.ChannelID] = append(s.messages[req.ChannelID], serviceMsg) - s.events[req.ChannelID] = append(s.events[req.ChannelID], serviceEvent) + s.appendChannelEventLocked(serviceEvent) s.updateForumTopicTopMessageLocked(req.ChannelID, serviceMsg) channel.TopMessageID = serviceMsg.ID channel.Pts = servicePts @@ -303,7 +303,7 @@ func (s *ChannelStore) UpdatePinnedMessage(_ context.Context, req domain.UpdateC SenderUserID: req.UserID, Pinned: req.Pinned, } - s.events[req.ChannelID] = append(s.events[req.ChannelID], event) + s.appendChannelEventLocked(event) logMsg := msg logMsg.Pinned = req.Pinned s.appendChannelAdminLogLocked(domain.ChannelAdminLogEvent{ @@ -372,7 +372,7 @@ func (s *ChannelStore) UnpinAllChannelMessages(_ context.Context, req domain.Unp SenderUserID: req.UserID, Pinned: false, } - s.events[req.ChannelID] = append(s.events[req.ChannelID], event) + s.appendChannelEventLocked(event) return domain.UpdateChannelPinnedMessageResult{ Channel: channel, Event: cloneChannelEvent(event), diff --git a/internal/store/memory/channel_message_send.go b/internal/store/memory/channel_message_send.go index 5604de10..41d64b2e 100644 --- a/internal/store/memory/channel_message_send.go +++ b/internal/store/memory/channel_message_send.go @@ -2,8 +2,10 @@ package memory import ( "context" + "fmt" "strings" "telesrv/internal/domain" + "telesrv/internal/store" "time" ) @@ -14,8 +16,27 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan if strings.TrimSpace(req.Message) == "" && req.Action == nil && req.Media.IsZero() && req.RichMessage.IsZero() { return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid } + var fingerprint []byte + var err error + if req.RandomID != 0 { + fingerprint, err = store.ChannelSendFingerprint(req) + if err != nil { + return domain.SendChannelMessageResult{}, err + } + req.IdempotencyFingerprint = fingerprint + } s.mu.Lock() defer s.mu.Unlock() + if req.RandomID != 0 { + if replay, found, replayErr := s.lookupChannelSendReplayLocked(domain.ChannelSendReplayRequest{ + ChannelID: req.ChannelID, + SenderUserID: req.UserID, + RandomID: req.RandomID, + IdempotencyFingerprint: fingerprint, + }); replayErr != nil || found { + return replay, replayErr + } + } channel, err := s.channelForMemberLocked(req.UserID, req.ChannelID) if err != nil { return domain.SendChannelMessageResult{}, err @@ -34,23 +55,6 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan if !canSendChannelMessageWithBoost(channel, member, fromBoostsApplied) { return domain.SendChannelMessageResult{}, domain.ErrChannelWriteForbidden } - if req.RandomID != 0 { - if id, ok := s.randomToID[channelRandomKey{channelID: req.ChannelID, userID: req.UserID, randomID: req.RandomID}]; ok { - msg, ok := s.findMessageLocked(req.ChannelID, id) - if ok { - event := s.eventForMessageLocked(req.ChannelID, id) - if event.Message.ID != 0 { - msg = event.Message - } - return domain.SendChannelMessageResult{ - Channel: channel, - Message: cloneChannelMessage(msg), - Event: event, - Duplicate: true, - }, nil - } - } - } if wait := channelSlowModeWait(channel, member, req.Date); wait > 0 { return domain.SendChannelMessageResult{}, domain.NewSlowModeWaitError(wait) } @@ -99,7 +103,7 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan Message: cloneChannelMessage(discussionMsg), } s.messages[linked.ID] = append(s.messages[linked.ID], discussionMsg) - s.events[linked.ID] = append(s.events[linked.ID], discussionEvent) + s.appendChannelEventLocked(discussionEvent) linked.TopMessageID = discussionMsgID linked.Pts = discussionPts s.channels[linked.ID] = linked @@ -145,6 +149,13 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan Pts: pts, } msg.Replies = s.channelMessageRepliesLocked(req.UserID, req.ChannelID, msg) + var sendSnapshot []byte + if req.RandomID != 0 { + sendSnapshot, err = store.EncodeChannelSendSnapshot(msg) + if err != nil { + return domain.SendChannelMessageResult{}, err + } + } event := domain.ChannelUpdateEvent{ ChannelID: req.ChannelID, Type: domain.ChannelUpdateNewMessage, @@ -155,7 +166,7 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan SenderUserID: req.UserID, } s.messages[req.ChannelID] = append(s.messages[req.ChannelID], msg) - s.events[req.ChannelID] = append(s.events[req.ChannelID], event) + s.appendChannelEventLocked(event) if !channel.Broadcast || channel.Megagroup { mentionTargets := req.MentionUserIDs if msg.ReplyTo != nil && msg.ReplyTo.MessageID > 0 { @@ -178,7 +189,11 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan }) } if req.RandomID != 0 { - s.randomToID[channelRandomKey{channelID: req.ChannelID, userID: req.UserID, randomID: req.RandomID}] = msg.ID + key := channelRandomKey{channelID: req.ChannelID, userID: req.UserID, randomID: req.RandomID} + s.randomToID[key] = msg.ID + replayKey := channelMessageReplayKey{channelID: req.ChannelID, messageID: msg.ID} + s.sendSnapshots[replayKey] = sendSnapshot + s.sendFingerprints[replayKey] = append([]byte(nil), fingerprint...) } channel.TopMessageID = msg.ID channel.Pts = pts @@ -209,6 +224,85 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan }, nil } +// LookupChannelSendReplay returns a committed regular-channel or monoforum send receipt without +// evaluating current membership, write permissions, slow mode or any message allocation path. +func (s *ChannelStore) LookupChannelSendReplay(_ context.Context, req domain.ChannelSendReplayRequest) (domain.SendChannelMessageResult, bool, error) { + if req.ChannelID == 0 || req.SenderUserID == 0 || req.RandomID == 0 { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("memory channel send replay: invalid scope") + } + if req.SavedPeer.ID == 0 { + if req.SavedPeer.Type != "" { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("memory channel send replay: incomplete saved peer scope") + } + } else if req.SavedPeer.Type != domain.PeerTypeUser { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("memory channel send replay: invalid saved peer scope") + } + if err := store.ValidateSendFingerprint(req.IdempotencyFingerprint, "channel send replay"); err != nil { + return domain.SendChannelMessageResult{}, false, err + } + s.mu.RLock() + defer s.mu.RUnlock() + return s.lookupChannelSendReplayLocked(req) +} + +func (s *ChannelStore) lookupChannelSendReplayLocked(req domain.ChannelSendReplayRequest) (domain.SendChannelMessageResult, bool, error) { + var id int + if req.SavedPeer.ID == 0 { + var found bool + id, found = s.randomToID[channelRandomKey{channelID: req.ChannelID, userID: req.SenderUserID, randomID: req.RandomID}] + if !found { + return domain.SendChannelMessageResult{}, false, nil + } + } else { + msg, found := s.findMonoforumDuplicateLocked(req.ChannelID, req.SenderUserID, req.SavedPeer, req.RandomID) + if !found { + return domain.SendChannelMessageResult{}, false, nil + } + id = msg.ID + } + replayKey := channelMessageReplayKey{channelID: req.ChannelID, messageID: id} + if !store.SameSendFingerprint(s.sendFingerprints[replayKey], req.IdempotencyFingerprint) { + return domain.SendChannelMessageResult{}, false, domain.ErrMessageRandomIDDuplicate + } + first, err := store.DecodeChannelSendSnapshot(s.sendSnapshots[replayKey]) + if err != nil { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("memory duplicate channel message snapshot: %w", err) + } + if first.ID != id || first.ChannelID != req.ChannelID || first.SenderUserID != req.SenderUserID || first.RandomID != req.RandomID || first.SavedPeer != req.SavedPeer { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("memory duplicate channel message snapshot disagrees with random_id receipt") + } + replay := first + var replayDelete *domain.ChannelUpdateEvent + if current, found := s.findMessageLocked(req.ChannelID, id); found && !current.Deleted { + replay = cloneChannelMessage(current) + } else if receipt := s.deleteReceipts[replayKey]; receipt != nil { + cloned := cloneChannelEvent(*receipt) + replayDelete = &cloned + } else { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("memory duplicate channel message %d is absent without a durable delete receipt", id) + } + channel, ok := s.channels[req.ChannelID] + if !ok { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("memory duplicate channel message %d has no channel", id) + } + event := domain.ChannelUpdateEvent{ + ChannelID: first.ChannelID, + Type: domain.ChannelUpdateNewMessage, + Pts: first.Pts, + PtsCount: 1, + Date: first.Date, + Message: cloneChannelMessage(replay), + SenderUserID: first.SenderUserID, + } + return domain.SendChannelMessageResult{ + Channel: cloneChannel(channel), + Message: cloneChannelMessage(replay), + Event: event, + Duplicate: true, + ReplayDeleteEvent: replayDelete, + }, true, nil +} + func channelDeliverySkipSet(ids []int64) map[int64]struct{} { if len(ids) == 0 { return nil @@ -267,7 +361,7 @@ func (s *ChannelStore) appendChannelServiceMessageLocked(channelID, senderUserID UserIDs: append([]int64(nil), action.UserIDs...), } s.messages[channelID] = append(s.messages[channelID], msg) - s.events[channelID] = append(s.events[channelID], event) + s.appendChannelEventLocked(event) return msg, event } diff --git a/internal/store/memory/channel_monoforum.go b/internal/store/memory/channel_monoforum.go index 70010d26..077f8190 100644 --- a/internal/store/memory/channel_monoforum.go +++ b/internal/store/memory/channel_monoforum.go @@ -7,6 +7,7 @@ import ( "time" "telesrv/internal/domain" + "telesrv/internal/store" ) // SendMonoforumMessage 向 monoforum(频道私信)虚拟频道发一条消息,按 saved_peer 分订阅者子会话。 @@ -16,8 +17,28 @@ func (s *ChannelStore) SendMonoforumMessage(_ context.Context, req domain.SendMo req.SavedPeer.Type != domain.PeerTypeUser || strings.TrimSpace(req.Message) == "" { return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid } + var fingerprint []byte + var err error + if req.RandomID != 0 { + fingerprint, err = store.MonoforumSendFingerprint(req) + if err != nil { + return domain.SendChannelMessageResult{}, err + } + req.IdempotencyFingerprint = fingerprint + } s.mu.Lock() defer s.mu.Unlock() + if req.RandomID != 0 { + if replay, found, replayErr := s.lookupChannelSendReplayLocked(domain.ChannelSendReplayRequest{ + ChannelID: req.MonoforumID, + SenderUserID: req.SenderUserID, + SavedPeer: req.SavedPeer, + RandomID: req.RandomID, + IdempotencyFingerprint: fingerprint, + }); replayErr != nil || found { + return replay, replayErr + } + } channel, ok := s.channels[req.MonoforumID] if !ok || channel.Deleted || !channel.Monoforum { return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid @@ -25,17 +46,6 @@ func (s *ChannelStore) SendMonoforumMessage(_ context.Context, req domain.SendMo if req.Date == 0 { req.Date = int(time.Now().Unix()) } - if req.RandomID != 0 { - // 去重维度 = (sender, saved_peer, random_id),与 postgres 迁移 0022 的唯一索引一致; - // 不复用账号级 randomToID(其按 channel+sender+random_id 三元组,会被跨子会话同 random_id 互相覆盖)。 - if dup, ok := s.findMonoforumDuplicateLocked(req.MonoforumID, req.SenderUserID, req.SavedPeer, req.RandomID); ok { - event := s.eventForMessageLocked(req.MonoforumID, dup.ID) - if event.Message.ID != 0 { - dup = event.Message - } - return domain.SendChannelMessageResult{Channel: cloneChannel(channel), Message: cloneChannelMessage(dup), Event: event, Duplicate: true}, nil - } - } pts := s.nextChannelPtsLocked(req.MonoforumID) msgID := s.nextChannelMessageIDLocked(req.MonoforumID) msg := domain.ChannelMessage{ @@ -50,6 +60,14 @@ func (s *ChannelStore) SendMonoforumMessage(_ context.Context, req domain.SendMo Entities: append([]domain.MessageEntity(nil), req.Entities...), Pts: pts, } + var sendSnapshot []byte + if req.RandomID != 0 { + var snapshotErr error + sendSnapshot, snapshotErr = store.EncodeChannelSendSnapshot(msg) + if snapshotErr != nil { + return domain.SendChannelMessageResult{}, snapshotErr + } + } event := domain.ChannelUpdateEvent{ ChannelID: req.MonoforumID, Type: domain.ChannelUpdateNewMessage, @@ -60,7 +78,12 @@ func (s *ChannelStore) SendMonoforumMessage(_ context.Context, req domain.SendMo SenderUserID: req.SenderUserID, } s.messages[req.MonoforumID] = append(s.messages[req.MonoforumID], msg) - s.events[req.MonoforumID] = append(s.events[req.MonoforumID], event) + if req.RandomID != 0 { + replayKey := channelMessageReplayKey{channelID: req.MonoforumID, messageID: msg.ID} + s.sendSnapshots[replayKey] = sendSnapshot + s.sendFingerprints[replayKey] = append([]byte(nil), fingerprint...) + } + s.appendChannelEventLocked(event) channel.TopMessageID = msgID channel.Pts = pts s.channels[req.MonoforumID] = channel @@ -75,7 +98,7 @@ func (s *ChannelStore) findMonoforumDuplicateLocked(monoforumID, senderUserID in msgs := s.messages[monoforumID] for i := len(msgs) - 1; i >= 0; i-- { m := msgs[i] - if !m.Deleted && m.RandomID == randomID && m.SenderUserID == senderUserID && m.SavedPeer == savedPeer { + if m.RandomID == randomID && m.SenderUserID == senderUserID && m.SavedPeer == savedPeer { return m, true } } diff --git a/internal/store/memory/channel_monoforum_send_test.go b/internal/store/memory/channel_monoforum_send_test.go index 0f2fe432..9652db64 100644 --- a/internal/store/memory/channel_monoforum_send_test.go +++ b/internal/store/memory/channel_monoforum_send_test.go @@ -2,6 +2,7 @@ package memory import ( "context" + "errors" "testing" "telesrv/internal/domain" @@ -69,6 +70,9 @@ func TestSendMonoforumMessageAndHistory(t *testing.T) { if !dup.Duplicate || dup.Message.ID != m1.Message.ID { t.Fatalf("dup = %+v, want duplicate of m1 id %d", dup.Message, m1.Message.ID) } + if _, err := store.SendMonoforumMessage(ctx, domain.SendMonoforumMessageRequest{MonoforumID: monoID, SenderUserID: 42, SavedPeer: sub, RandomID: 111, Message: "changed", Date: 1_700_001_004}); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("conflicting monoforum replay err=%v, want ErrMessageRandomIDDuplicate", err) + } hist, err := store.ListMonoforumHistory(ctx, domain.MonoforumHistoryFilter{MonoforumID: monoID, SavedPeer: sub, Limit: 10}) if err != nil { @@ -131,4 +135,21 @@ func TestSendMonoforumMessageAndHistory(t *testing.T) { if dialogs.Dialogs[1].SavedPeer != sub || dialogs.Dialogs[1].TopMessageID == 0 { t.Fatalf("dialogs[1] = %+v, want sub with top message", dialogs.Dialogs[1]) } + store.mu.Lock() + _, deleteEvent, _, err := store.deleteChannelMessagesLocked(store.channels[monoID], domain.ChannelMember{ChannelID: monoID, UserID: 1, Role: domain.ChannelRoleCreator, Status: domain.ChannelMemberActive}, []int{a.Message.ID}, 1, 1_700_001_013) + store.mu.Unlock() + if err != nil { + t.Fatalf("delete monoforum message: %v", err) + } + ptsBeforeReplay, eventsBeforeReplay := store.ptsSeq[monoID], len(store.events[monoID]) + deletedReplay, err := store.SendMonoforumMessage(ctx, domain.SendMonoforumMessageRequest{MonoforumID: monoID, SenderUserID: 1, SavedPeer: sub, RandomID: 9001, Message: "to sub", Date: 1_700_001_014}) + if err != nil { + t.Fatalf("replay deleted monoforum message: %v", err) + } + if !deletedReplay.Duplicate || deletedReplay.Message.ID != a.Message.ID || deletedReplay.Message.Body != "to sub" || deletedReplay.ReplayDeleteEvent == nil || deletedReplay.ReplayDeleteEvent.Pts != deleteEvent.Pts { + t.Fatalf("deleted monoforum replay = %+v, want first snapshot + durable delete %+v", deletedReplay, deleteEvent) + } + if store.ptsSeq[monoID] != ptsBeforeReplay || len(store.events[monoID]) != eventsBeforeReplay { + t.Fatalf("deleted monoforum replay mutated pts/events = %d/%d, want %d/%d", store.ptsSeq[monoID], len(store.events[monoID]), ptsBeforeReplay, eventsBeforeReplay) + } } diff --git a/internal/store/memory/channel_recovery_test.go b/internal/store/memory/channel_recovery_test.go new file mode 100644 index 00000000..2ff89d42 --- /dev/null +++ b/internal/store/memory/channel_recovery_test.go @@ -0,0 +1,22 @@ +package memory + +import ( + "context" + "testing" +) + +func TestMaxChannelPtsBatchUsesOneSnapshotAndOmitsMissing(t *testing.T) { + channels := NewChannelStore() + channels.mu.Lock() + channels.ptsSeq[10] = 7 + channels.ptsSeq[20] = 11 + channels.mu.Unlock() + + got, err := channels.MaxChannelPtsBatch(context.Background(), []int64{20, 999, 10, 20}) + if err != nil { + t.Fatalf("MaxChannelPtsBatch: %v", err) + } + if len(got) != 2 || got[10] != 7 || got[20] != 11 { + t.Fatalf("batch pts = %v, want map[10:7 20:11] with missing id omitted", got) + } +} diff --git a/internal/store/memory/channel_send_idempotency_test.go b/internal/store/memory/channel_send_idempotency_test.go new file mode 100644 index 00000000..a2bbe10b --- /dev/null +++ b/internal/store/memory/channel_send_idempotency_test.go @@ -0,0 +1,76 @@ +package memory + +import ( + "context" + "errors" + "testing" + + "telesrv/internal/domain" +) + +func TestChannelRandomIDReplayUsesCurrentSnapshotAndDurableDeleteMemory(t *testing.T) { + ctx := context.Background() + channels := NewChannelStore() + created, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: 1, + Title: "replay convergence", + Megagroup: true, + Date: 1_700_001_000, + }) + if err != nil { + t.Fatalf("create channel: %v", err) + } + req := domain.SendChannelMessageRequest{ + UserID: 1, ChannelID: created.Channel.ID, RandomID: 77001, + Message: "original", Date: 1_700_001_001, + } + first, err := channels.SendChannelMessage(ctx, req) + if err != nil { + t.Fatalf("send channel message: %v", err) + } + conflict := req + conflict.Message = "same random id, different intent" + if _, err := channels.SendChannelMessage(ctx, conflict); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("conflicting channel replay err=%v, want ErrMessageRandomIDDuplicate", err) + } + edited, err := channels.EditChannelMessage(ctx, domain.EditChannelMessageRequest{ + UserID: 1, ChannelID: created.Channel.ID, ID: first.Message.ID, + Message: "edited", EditDate: 1_700_001_002, + }) + if err != nil { + t.Fatalf("edit channel message: %v", err) + } + ptsBeforeReplay := channels.ptsSeq[created.Channel.ID] + eventsBeforeReplay := len(channels.events[created.Channel.ID]) + replay, err := channels.SendChannelMessage(ctx, req) + if err != nil { + t.Fatalf("replay after edit: %v", err) + } + if !replay.Duplicate || replay.Message.Body != "edited" || replay.Message.Pts != edited.Message.Pts || replay.Event.Pts != first.Event.Pts || replay.ReplayDeleteEvent != nil { + t.Fatalf("replay after edit = %+v, want current snapshot with first-send pts", replay) + } + if channels.ptsSeq[created.Channel.ID] != ptsBeforeReplay || len(channels.events[created.Channel.ID]) != eventsBeforeReplay { + t.Fatalf("edit replay mutated channel pts/events = %d/%d, want %d/%d", channels.ptsSeq[created.Channel.ID], len(channels.events[created.Channel.ID]), ptsBeforeReplay, eventsBeforeReplay) + } + deleted, err := channels.DeleteChannelMessages(ctx, domain.DeleteChannelMessagesRequest{ + UserID: 1, ChannelID: created.Channel.ID, IDs: []int{first.Message.ID}, Date: 1_700_001_003, + }) + if err != nil { + t.Fatalf("delete channel message: %v", err) + } + ptsBeforeReplay = channels.ptsSeq[created.Channel.ID] + eventsBeforeReplay = len(channels.events[created.Channel.ID]) + replay, err = channels.SendChannelMessage(ctx, req) + if err != nil { + t.Fatalf("replay after delete: %v", err) + } + if !replay.Duplicate || replay.Message.Body != "original" || replay.Message.Pts != first.Message.Pts || replay.Event.Pts != first.Event.Pts { + t.Fatalf("replay after delete = %+v, want immutable first snapshot", replay) + } + if replay.ReplayDeleteEvent == nil || replay.ReplayDeleteEvent.Pts != deleted.Event.Pts || len(replay.ReplayDeleteEvent.MessageIDs) != 1 || replay.ReplayDeleteEvent.MessageIDs[0] != first.Message.ID { + t.Fatalf("replay delete = %+v, want durable event %+v", replay.ReplayDeleteEvent, deleted.Event) + } + if channels.ptsSeq[created.Channel.ID] != ptsBeforeReplay || len(channels.events[created.Channel.ID]) != eventsBeforeReplay { + t.Fatalf("delete replay mutated channel pts/events = %d/%d, want %d/%d", channels.ptsSeq[created.Channel.ID], len(channels.events[created.Channel.ID]), ptsBeforeReplay, eventsBeforeReplay) + } +} diff --git a/internal/store/memory/channel_store.go b/internal/store/memory/channel_store.go index af8fbae7..fac32d7a 100644 --- a/internal/store/memory/channel_store.go +++ b/internal/store/memory/channel_store.go @@ -13,6 +13,11 @@ type channelRandomKey struct { randomID int64 } +type channelMessageReplayKey struct { + channelID int64 + messageID int +} + type boostSlotKey struct { userID int64 slot int @@ -58,33 +63,37 @@ func (w channelReadWatermark) advance(userID int64, maxID int) channelReadWaterm // ChannelStore is an in-memory channel/supergroup store for tests and local development. type ChannelStore struct { - mu sync.RWMutex - nextID int64 - nextHash int64 - channels map[int64]domain.Channel - members map[int64]map[int64]domain.ChannelMember - dialogs map[int64]map[int64]domain.ChannelDialog - topics map[int64]map[int]domain.ChannelForumTopic - messages map[int64][]domain.ChannelMessage - reactions map[int64]map[int]map[int64][]domain.ChannelMessagePeerReaction + mu sync.RWMutex + nextID int64 + nextHash int64 + channels map[int64]domain.Channel + members map[int64]map[int64]domain.ChannelMember + dialogs map[int64]map[int64]domain.ChannelDialog + topics map[int64]map[int]domain.ChannelForumTopic + messages map[int64][]domain.ChannelMessage + reactions map[int64]map[int]map[int64][]domain.ChannelMessagePeerReaction // paidReactions 是 per-(channel,message,user) 付费 reaction 累计星数 + 匿名标志。 - paidReactions map[int64]map[int]map[int64]memoryPaidReaction - top map[int64]map[string]domain.TopMessageReaction - recent map[int64]map[string]domain.RecentMessageReaction - savedTags map[int64]map[string]domain.SavedReactionTag - mentions map[int64]map[int64]map[int]memoryMention - msgViews map[int64]map[int]int - msgViewers map[int64]map[int]map[int64]struct{} - events map[int64][]domain.ChannelUpdateEvent - adminLogs map[int64][]domain.ChannelAdminLogEvent - invites map[string]domain.ChannelInvite - importers map[int64]map[int64]domain.ChannelInviteImporter - msgSeq map[int64]int - ptsSeq map[int64]int - logSeq map[int64]int64 - randomToID map[channelRandomKey]int - boostSlots map[boostSlotKey]domain.PremiumBoostSlot - readMarks map[int64]channelReadWatermark + paidReactions map[int64]map[int]map[int64]memoryPaidReaction + top map[int64]map[string]domain.TopMessageReaction + recent map[int64]map[string]domain.RecentMessageReaction + savedTags map[int64]map[string]domain.SavedReactionTag + mentions map[int64]map[int64]map[int]memoryMention + msgViews map[int64]map[int]int + msgViewers map[int64]map[int]map[int64]struct{} + events map[int64][]domain.ChannelUpdateEvent + retention map[int64]domain.ChannelUpdateRetentionCheckpoint + adminLogs map[int64][]domain.ChannelAdminLogEvent + invites map[string]domain.ChannelInvite + importers map[int64]map[int64]domain.ChannelInviteImporter + msgSeq map[int64]int + ptsSeq map[int64]int + logSeq map[int64]int64 + randomToID map[channelRandomKey]int + sendSnapshots map[channelMessageReplayKey][]byte + sendFingerprints map[channelMessageReplayKey][]byte + deleteReceipts map[channelMessageReplayKey]*domain.ChannelUpdateEvent + boostSlots map[boostSlotKey]domain.PremiumBoostSlot + readMarks map[int64]channelReadWatermark // topicReads 是 per-(channel,user,topic) 已读水位（forum 话题独立已读，不碰频道级 member 水位）。 topicReads map[int64]map[int64]map[int]memoryTopicRead // polls 是共享 poll 权威（与 MessageStore 同一实例）；nil 时 poll 链路按未接入处理。 @@ -99,31 +108,35 @@ func (s *ChannelStore) AttachPollStore(polls *PollStore) { // NewChannelStore creates an in-memory ChannelStore. func NewChannelStore() *ChannelStore { return &ChannelStore{ - nextID: firstMemoryChannelID, - nextHash: 900000000000, - channels: make(map[int64]domain.Channel), - members: make(map[int64]map[int64]domain.ChannelMember), - dialogs: make(map[int64]map[int64]domain.ChannelDialog), - topics: make(map[int64]map[int]domain.ChannelForumTopic), - messages: make(map[int64][]domain.ChannelMessage), - reactions: make(map[int64]map[int]map[int64][]domain.ChannelMessagePeerReaction), - paidReactions: make(map[int64]map[int]map[int64]memoryPaidReaction), - top: make(map[int64]map[string]domain.TopMessageReaction), - recent: make(map[int64]map[string]domain.RecentMessageReaction), - savedTags: make(map[int64]map[string]domain.SavedReactionTag), - mentions: make(map[int64]map[int64]map[int]memoryMention), - msgViews: make(map[int64]map[int]int), - msgViewers: make(map[int64]map[int]map[int64]struct{}), - events: make(map[int64][]domain.ChannelUpdateEvent), - adminLogs: make(map[int64][]domain.ChannelAdminLogEvent), - invites: make(map[string]domain.ChannelInvite), - importers: make(map[int64]map[int64]domain.ChannelInviteImporter), - msgSeq: make(map[int64]int), - ptsSeq: make(map[int64]int), - logSeq: make(map[int64]int64), - randomToID: make(map[channelRandomKey]int), - boostSlots: make(map[boostSlotKey]domain.PremiumBoostSlot), - readMarks: make(map[int64]channelReadWatermark), - topicReads: make(map[int64]map[int64]map[int]memoryTopicRead), + nextID: firstMemoryChannelID, + nextHash: 900000000000, + channels: make(map[int64]domain.Channel), + members: make(map[int64]map[int64]domain.ChannelMember), + dialogs: make(map[int64]map[int64]domain.ChannelDialog), + topics: make(map[int64]map[int]domain.ChannelForumTopic), + messages: make(map[int64][]domain.ChannelMessage), + reactions: make(map[int64]map[int]map[int64][]domain.ChannelMessagePeerReaction), + paidReactions: make(map[int64]map[int]map[int64]memoryPaidReaction), + top: make(map[int64]map[string]domain.TopMessageReaction), + recent: make(map[int64]map[string]domain.RecentMessageReaction), + savedTags: make(map[int64]map[string]domain.SavedReactionTag), + mentions: make(map[int64]map[int64]map[int]memoryMention), + msgViews: make(map[int64]map[int]int), + msgViewers: make(map[int64]map[int]map[int64]struct{}), + events: make(map[int64][]domain.ChannelUpdateEvent), + retention: make(map[int64]domain.ChannelUpdateRetentionCheckpoint), + adminLogs: make(map[int64][]domain.ChannelAdminLogEvent), + invites: make(map[string]domain.ChannelInvite), + importers: make(map[int64]map[int64]domain.ChannelInviteImporter), + msgSeq: make(map[int64]int), + ptsSeq: make(map[int64]int), + logSeq: make(map[int64]int64), + randomToID: make(map[channelRandomKey]int), + sendSnapshots: make(map[channelMessageReplayKey][]byte), + sendFingerprints: make(map[channelMessageReplayKey][]byte), + deleteReceipts: make(map[channelMessageReplayKey]*domain.ChannelUpdateEvent), + boostSlots: make(map[boostSlotKey]domain.PremiumBoostSlot), + readMarks: make(map[int64]channelReadWatermark), + topicReads: make(map[int64]map[int64]map[int]memoryTopicRead), } } diff --git a/internal/store/memory/channel_test.go b/internal/store/memory/channel_test.go index d738ced1..f1455157 100644 --- a/internal/store/memory/channel_test.go +++ b/internal/store/memory/channel_test.go @@ -1093,8 +1093,12 @@ func TestChannelMessageReplyMarkupSurvivesReadPaths(t *testing.T) { UserID: 1, ChannelID: created.Channel.ID, RandomID: 38_001, - Message: "duplicate must not replace", - Date: 1_700_000_382, + Message: "via inline keyboard", + ViaBotID: 99, + ReplyMarkup: &domain.MessageReplyMarkup{Inline: [][]domain.MarkupButton{{ + {Type: domain.MarkupButtonCallback, Text: "Open", Data: []byte{0x00, 0xff, 0x42}}, + }}}, + Date: 1_700_000_382, }) if err != nil { t.Fatalf("duplicate send: %v", err) diff --git a/internal/store/memory/channel_update_retention_test.go b/internal/store/memory/channel_update_retention_test.go new file mode 100644 index 00000000..fc9ae772 --- /dev/null +++ b/internal/store/memory/channel_update_retention_test.go @@ -0,0 +1,150 @@ +package memory + +import ( + "context" + "strings" + "testing" + "time" + + "telesrv/internal/domain" +) + +func TestChannelUpdateRetentionFloorDifferenceAndDirtyCheckpointMemory(t *testing.T) { + ctx := context.Background() + store := NewChannelStore() + const ownerID int64 = 701 + created, err := store.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: ownerID, + Title: "retention memory", + Megagroup: true, + Date: 1_700_010_000, + }) + if err != nil { + t.Fatalf("create channel: %v", err) + } + channelID := created.Channel.ID + sent := make([]domain.SendChannelMessageResult, 0, 3) + for i := 1; i <= 3; i++ { + result, err := store.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + UserID: ownerID, + ChannelID: channelID, + RandomID: int64(9000 + i), + Message: "retention", + Date: 1_700_010_000 + i, + }) + if err != nil { + t.Fatalf("send message %d: %v", i, err) + } + sent = append(sent, result) + } + + // Delete create + first two messages. The third message remains as the normal incremental page. + pruned, err := store.PruneChannelUpdateEvents(ctx, channelID, sent[1].Event.Pts, 100) + if err != nil { + t.Fatalf("prune channel updates: %v", err) + } + if pruned.Deleted != 3 || pruned.Checkpoint.RetainedThroughPts != sent[1].Event.Pts { + t.Fatalf("prune result = %+v, want deleted=3 floor=%d", pruned, sent[1].Event.Pts) + } + if pruned.Checkpoint.LatestPts != sent[2].Event.Pts || pruned.Checkpoint.LatestEventDate != sent[2].Event.Date { + t.Fatalf("checkpoint latest = %+v, want pts/date %d/%d", pruned.Checkpoint, sent[2].Event.Pts, sent[2].Event.Date) + } + + below, err := store.ListChannelDifference(ctx, domain.ChannelDifferenceRequest{ + UserID: ownerID, ChannelID: channelID, Pts: pruned.Checkpoint.RetainedThroughPts - 1, Limit: 100, + }) + if err != nil { + t.Fatalf("difference below retained floor: %v", err) + } + if !below.TooLong || below.Pts != sent[2].Event.Pts || below.Dialog.ChannelID != channelID { + t.Fatalf("difference below floor = %+v, want complete too-long snapshot at pts %d", below, sent[2].Event.Pts) + } + atFloor, err := store.ListChannelDifference(ctx, domain.ChannelDifferenceRequest{ + UserID: ownerID, ChannelID: channelID, Pts: pruned.Checkpoint.RetainedThroughPts, Limit: 100, + }) + if err != nil { + t.Fatalf("difference at retained floor: %v", err) + } + if atFloor.TooLong || len(atFloor.Events) != 1 || atFloor.Events[0].Pts != sent[2].Event.Pts { + t.Fatalf("difference at floor = %+v, want one normal incremental event at pts %d", atFloor, sent[2].Event.Pts) + } + + // Remove the remaining row. Dirty-channel recovery must still use checkpoint.latest_event_date. + allPruned, err := store.PruneChannelUpdateEvents(ctx, channelID, sent[2].Event.Pts, 100) + if err != nil { + t.Fatalf("prune remaining channel updates: %v", err) + } + if allPruned.Deleted != 1 || len(store.events[channelID]) != 0 { + t.Fatalf("remaining prune = %+v events=%v, want empty event log", allPruned, store.events[channelID]) + } + dirty, err := store.ListDirtyActiveChannelsForUser(ctx, ownerID, sent[2].Event.Date-1, 0, 10) + if err != nil { + t.Fatalf("list dirty channels after prune: %v", err) + } + if len(dirty) != 1 || dirty[0].ChannelID != channelID || dirty[0].Pts != sent[2].Event.Pts { + t.Fatalf("dirty channels after prune = %+v, want channel %d pts %d", dirty, channelID, sent[2].Event.Pts) + } +} + +func TestPruneChannelUpdateEventsRejectsInvalidPtsCountMemory(t *testing.T) { + ctx := context.Background() + store := NewChannelStore() + created, err := store.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: 901, + Title: "invalid retention event", + Megagroup: true, + Date: 1_700_020_000, + }) + if err != nil { + t.Fatalf("create channel: %v", err) + } + channelID := created.Channel.ID + channel := store.channels[channelID] + channel.Pts++ + store.channels[channelID] = channel + store.ptsSeq[channelID] = channel.Pts + store.appendChannelEventLocked(domain.ChannelUpdateEvent{ + ChannelID: channelID, + Type: domain.ChannelUpdateNoop, + Pts: channel.Pts, + PtsCount: 0, + Date: 1_700_020_001, + }) + + _, err = store.PruneChannelUpdateEvents(ctx, channelID, channel.Pts, 100) + if err == nil || !strings.Contains(err.Error(), "invalid pts_count=0") { + t.Fatalf("prune invalid event err = %v, want fail-fast pts_count error", err) + } +} + +func TestDeleteExpiredChannelUpdateEventsIsBoundedMemory(t *testing.T) { + ctx := context.Background() + store := NewChannelStore() + created, err := store.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: 801, + Title: "expired retention memory", + Megagroup: true, + Date: 1_600_000_000, + }) + if err != nil { + t.Fatalf("create channel: %v", err) + } + for i := 1; i <= 3; i++ { + if _, err := store.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + UserID: 801, ChannelID: created.Channel.ID, RandomID: int64(8000 + i), Message: "old", Date: 1_600_000_000 + i, + }); err != nil { + t.Fatalf("send old message %d: %v", i, err) + } + } + deleted, err := store.DeleteExpiredChannelUpdateEvents(ctx, time.Hour, 2) + if err != nil { + t.Fatalf("delete expired channel updates: %v", err) + } + if deleted != 2 { + t.Fatalf("deleted = %d, want bounded batch 2", deleted) + } + checkpoint := store.retention[created.Channel.ID] + if checkpoint.RetainedThroughPts != 2 || len(store.events[created.Channel.ID]) != 2 { + t.Fatalf("after bounded prune checkpoint=%+v events=%d, want floor=2 and 2 rows", checkpoint, len(store.events[created.Channel.ID])) + } +} diff --git a/internal/store/memory/channel_updates.go b/internal/store/memory/channel_updates.go index 33e9900b..3776ac7c 100644 --- a/internal/store/memory/channel_updates.go +++ b/internal/store/memory/channel_updates.go @@ -2,7 +2,11 @@ package memory import ( "context" + "fmt" + "sort" "strings" + "time" + "telesrv/internal/domain" ) @@ -37,7 +41,8 @@ func (s *ChannelStore) ListChannelDifference(_ context.Context, req domain.Chann Dialog: dialog, }, nil } - if channel.Pts-req.Pts > limit { + checkpoint := s.channelUpdateCheckpointLocked(req.ChannelID, channel) + if req.Pts < checkpoint.RetainedThroughPts || channel.Pts-req.Pts > limit { messages := make([]domain.ChannelMessage, 0, domain.MaxChannelDifferenceTooLongMessages) for i := len(s.messages[req.ChannelID]) - 1; i >= 0 && len(messages) < domain.MaxChannelDifferenceTooLongMessages; i-- { msg := s.messages[req.ChannelID][i] @@ -122,6 +127,170 @@ func (s *ChannelStore) MaxChannelPts(_ context.Context, channelID int64) (int, e return s.ptsSeq[channelID], nil } +func (s *ChannelStore) MaxChannelPtsBatch(_ context.Context, channelIDs []int64) (map[int64]int, error) { + out := make(map[int64]int, len(channelIDs)) + s.mu.RLock() + for _, channelID := range channelIDs { + if pts, ok := s.ptsSeq[channelID]; ok { + out[channelID] = pts + } + } + s.mu.RUnlock() + return out, nil +} + +// appendChannelEventLocked is the only memory-store append boundary for channel-scoped durable +// events. Keeping the checkpoint current here mirrors the PostgreSQL event+checkpoint transaction. +func (s *ChannelStore) appendChannelEventLocked(event domain.ChannelUpdateEvent) { + s.events[event.ChannelID] = append(s.events[event.ChannelID], event) + checkpoint := s.retention[event.ChannelID] + checkpoint.ChannelID = event.ChannelID + if event.Pts > checkpoint.LatestPts { + checkpoint.LatestPts = event.Pts + } + if event.Date > checkpoint.LatestEventDate { + checkpoint.LatestEventDate = event.Date + } + s.retention[event.ChannelID] = checkpoint +} + +func (s *ChannelStore) channelUpdateCheckpointLocked(channelID int64, channel domain.Channel) domain.ChannelUpdateRetentionCheckpoint { + checkpoint := s.retention[channelID] + checkpoint.ChannelID = channelID + if channel.Pts > checkpoint.LatestPts { + checkpoint.LatestPts = channel.Pts + } + for _, event := range s.events[channelID] { + if event.Pts > checkpoint.LatestPts { + checkpoint.LatestPts = event.Pts + } + if event.Date > checkpoint.LatestEventDate { + checkpoint.LatestEventDate = event.Date + } + } + return checkpoint +} + +// PruneChannelUpdateEvents removes a bounded contiguous prefix and advances the retained floor in +// the same memory-store critical section. throughPts may land inside a pts_count interval; that row +// is retained because channel event rows are indivisible. +func (s *ChannelStore) PruneChannelUpdateEvents(_ context.Context, channelID int64, throughPts, limit int) (domain.ChannelUpdateRetentionResult, error) { + s.mu.Lock() + defer s.mu.Unlock() + return s.pruneChannelUpdateEventsLocked(channelID, throughPts, 0, limit) +} + +// DeleteExpiredChannelUpdateEvents uses the oldest retained event of each channel as an indexed-seek +// analogue, then prunes candidates oldest-first. There is no offset scan and total deleted rows never +// exceeds limit. +func (s *ChannelStore) DeleteExpiredChannelUpdateEvents(_ context.Context, olderThan time.Duration, limit int) (int, error) { + if olderThan <= 0 { + return 0, nil + } + limit = normalizeChannelRetentionLimit(limit) + cutoff := int(time.Now().Add(-olderThan).Unix()) + + s.mu.Lock() + defer s.mu.Unlock() + type candidate struct { + channelID int64 + date int + } + candidates := make([]candidate, 0) + for channelID, channel := range s.channels { + checkpoint := s.channelUpdateCheckpointLocked(channelID, channel) + for _, event := range s.events[channelID] { + if event.Pts <= checkpoint.RetainedThroughPts { + continue + } + if event.Date < cutoff { + candidates = append(candidates, candidate{channelID: channelID, date: event.Date}) + } + break + } + } + sort.Slice(candidates, func(i, j int) bool { + if candidates[i].date == candidates[j].date { + return candidates[i].channelID < candidates[j].channelID + } + return candidates[i].date < candidates[j].date + }) + + deleted := 0 + for _, item := range candidates { + if deleted >= limit { + break + } + channel := s.channels[item.channelID] + result, err := s.pruneChannelUpdateEventsLocked(item.channelID, channel.Pts, cutoff, limit-deleted) + if err != nil { + return deleted, err + } + deleted += result.Deleted + } + return deleted, nil +} + +func (s *ChannelStore) pruneChannelUpdateEventsLocked(channelID int64, throughPts, beforeDate, limit int) (domain.ChannelUpdateRetentionResult, error) { + channel, ok := s.channels[channelID] + if !ok || channelID == 0 || throughPts < 0 { + return domain.ChannelUpdateRetentionResult{}, domain.ErrChannelInvalid + } + limit = normalizeChannelRetentionLimit(limit) + checkpoint := s.channelUpdateCheckpointLocked(channelID, channel) + if throughPts > checkpoint.LatestPts { + throughPts = checkpoint.LatestPts + } + if throughPts <= checkpoint.RetainedThroughPts { + s.retention[channelID] = checkpoint + return domain.ChannelUpdateRetentionResult{Checkpoint: checkpoint}, nil + } + + cursor := checkpoint.RetainedThroughPts + deleted := 0 + keep := make([]domain.ChannelUpdateEvent, 0, len(s.events[channelID])) + canPrune := true + for _, event := range s.events[channelID] { + if event.Pts <= checkpoint.RetainedThroughPts { + keep = append(keep, event) + continue + } + if !canPrune || deleted >= limit || event.Pts > throughPts || (beforeDate > 0 && event.Date >= beforeDate) { + canPrune = false + keep = append(keep, event) + continue + } + ptsCount := event.PtsCount + if ptsCount <= 0 { + return domain.ChannelUpdateRetentionResult{}, fmt.Errorf( + "prune channel update events: channel %d has invalid pts_count=%d at pts=%d", + channelID, ptsCount, event.Pts, + ) + } + if event.Pts != cursor+ptsCount { + return domain.ChannelUpdateRetentionResult{}, fmt.Errorf( + "prune channel update events: channel %d has gap after pts %d: event pts=%d pts_count=%d", + channelID, cursor, event.Pts, ptsCount, + ) + } + cursor = event.Pts + deleted++ + } + if deleted > 0 { + s.events[channelID] = keep + checkpoint.RetainedThroughPts = cursor + } + s.retention[channelID] = checkpoint + return domain.ChannelUpdateRetentionResult{Checkpoint: checkpoint, Deleted: deleted}, nil +} + +func normalizeChannelRetentionLimit(limit int) int { + if limit <= 0 || limit > domain.MaxChannelUpdateRetentionBatch { + return domain.MaxChannelUpdateRetentionBatch + } + return limit +} + func (s *ChannelStore) nextChannelPtsLocked(channelID int64) int { s.ptsSeq[channelID]++ return s.ptsSeq[channelID] diff --git a/internal/store/memory/code_cas.go b/internal/store/memory/code_cas.go new file mode 100644 index 00000000..31c79528 --- /dev/null +++ b/internal/store/memory/code_cas.go @@ -0,0 +1,70 @@ +package memory + +import ( + "context" + + "telesrv/internal/store" +) + +func (s *CodeStore) GetSnapshot(_ context.Context, hash string) (store.PhoneCodeSnapshot, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + entry, found := s.liveCodeLocked(hash) + if !found { + return store.PhoneCodeSnapshot{}, false, nil + } + if entry.code.Version != store.PhoneCodeVersionCurrent || entry.code.Revision == "" { + s.deleteCodeLocked(hash, entry.code) + return store.PhoneCodeSnapshot{}, false, nil + } + return store.PhoneCodeSnapshot{Record: entry.code, Revision: entry.code.Revision}, true, nil +} + +func (s *CodeStore) CompareAndUpdate(_ context.Context, hash, expectedRevision string, next store.PhoneCode) (bool, error) { + if expectedRevision == "" || next.Version != store.PhoneCodeVersionCurrent || next.Purpose != "" { + return false, nil + } + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return false, err + } + next.Revision = revision + + s.mu.Lock() + defer s.mu.Unlock() + entry, found := s.liveCodeLocked(hash) + if !found { + return false, nil + } + if entry.code.Version != store.PhoneCodeVersionCurrent || entry.code.Revision == "" { + s.deleteCodeLocked(hash, entry.code) + return false, nil + } + if entry.code.Purpose != "" || entry.code.Revision != expectedRevision { + return false, nil + } + entry.code = next + s.m[hash] = entry + return true, nil +} + +func (s *CodeStore) CompareAndDelete(_ context.Context, hash, expectedRevision string) (bool, error) { + if expectedRevision == "" { + return false, nil + } + s.mu.Lock() + defer s.mu.Unlock() + entry, found := s.liveCodeLocked(hash) + if !found { + return false, nil + } + if entry.code.Version != store.PhoneCodeVersionCurrent || entry.code.Revision == "" { + s.deleteCodeLocked(hash, entry.code) + return false, nil + } + if entry.code.Purpose != "" || entry.code.Revision != expectedRevision { + return false, nil + } + s.deleteCodeLocked(hash, entry.code) + return true, nil +} diff --git a/internal/store/memory/code_cas_test.go b/internal/store/memory/code_cas_test.go new file mode 100644 index 00000000..1828533a --- /dev/null +++ b/internal/store/memory/code_cas_test.go @@ -0,0 +1,213 @@ +package memory + +import ( + "context" + "fmt" + "sync" + "testing" + "time" + + "telesrv/internal/store" +) + +func TestCodeStoreRevisionCAS(t *testing.T) { + ctx := context.Background() + codes := NewCodeStore() + record := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016201", + Code: "111111", + Channel: "email_setup", + PendingEmail: "first@example.test", + MaxAttempts: 5, + } + if err := codes.Set(ctx, "email-fixed", record, time.Minute); err != nil { + t.Fatal(err) + } + snapshot, found, err := codes.GetSnapshot(ctx, "email-fixed") + if err != nil || !found || snapshot.Revision == "" || snapshot.Record.Revision != snapshot.Revision { + t.Fatalf("snapshot=%+v found=%v err=%v", snapshot, found, err) + } + originalExpiry := codes.m["email-fixed"].expires + + next := snapshot.Record + next.Code = "222222" + next.Attempts = 1 + if applied, err := codes.CompareAndUpdate(ctx, "email-fixed", "stale-token", next); err != nil || applied { + t.Fatalf("wrong-token update applied=%v err=%v", applied, err) + } + unchanged, found, err := codes.GetSnapshot(ctx, "email-fixed") + if err != nil || !found || unchanged.Revision != snapshot.Revision || unchanged.Record.Code != record.Code { + t.Fatalf("after wrong-token snapshot=%+v found=%v err=%v", unchanged, found, err) + } + if applied, err := codes.CompareAndUpdate(ctx, "email-fixed", snapshot.Revision, next); err != nil || !applied { + t.Fatalf("current-token update applied=%v err=%v", applied, err) + } + updated, found, err := codes.GetSnapshot(ctx, "email-fixed") + if err != nil || !found || updated.Record.Code != next.Code || updated.Record.Attempts != 1 || updated.Revision == snapshot.Revision { + t.Fatalf("updated snapshot=%+v found=%v err=%v", updated, found, err) + } + if expiry := codes.m["email-fixed"].expires; !expiry.Equal(originalExpiry) { + t.Fatalf("CAS update expiry=%v, want unchanged %v", expiry, originalExpiry) + } + if applied, err := codes.CompareAndDelete(ctx, "email-fixed", snapshot.Revision); err != nil || applied { + t.Fatalf("stale delete applied=%v err=%v", applied, err) + } + if applied, err := codes.CompareAndDelete(ctx, "email-fixed", updated.Revision); err != nil || !applied { + t.Fatalf("current delete applied=%v err=%v", applied, err) + } + if _, found, _ := codes.GetSnapshot(ctx, "email-fixed"); found { + t.Fatal("CAS-deleted code remains") + } +} + +func TestCodeStoreRevisionCASFailClosedAndScopeIsolation(t *testing.T) { + ctx := context.Background() + codes := NewCodeStore() + now := time.Now().Add(time.Minute) + codes.m["legacy"] = codeEntry{ + code: store.PhoneCode{Version: 0, Phone: "15550016202", Code: "12345"}, + expires: now, + } + if _, found, err := codes.GetSnapshot(ctx, "legacy"); err != nil || found { + t.Fatalf("legacy snapshot found=%v err=%v", found, err) + } + if _, found := codes.m["legacy"]; found { + t.Fatal("legacy snapshot record was not deleted") + } + codes.m["no-revision"] = codeEntry{ + code: store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016202", + Code: "12345", + }, + expires: now, + } + if _, found, err := codes.GetSnapshot(ctx, "no-revision"); err != nil || found { + t.Fatalf("revisionless snapshot found=%v err=%v", found, err) + } + if _, found := codes.m["no-revision"]; found { + t.Fatal("revisionless snapshot record was not deleted") + } + + scoped := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016203", + Code: "12345", + Purpose: store.PhoneCodePurposeChangePhone, + UserID: 42, + AuthKeyID: [8]byte{1}, + } + if err := codes.Set(ctx, "scoped-cas", scoped, time.Minute); err != nil { + t.Fatal(err) + } + snapshot, found, err := codes.GetSnapshot(ctx, "scoped-cas") + if err != nil || !found { + t.Fatalf("scoped snapshot found=%v err=%v", found, err) + } + if applied, err := codes.CompareAndUpdate(ctx, "scoped-cas", snapshot.Revision, snapshot.Record); err != nil || applied { + t.Fatalf("scoped update applied=%v err=%v", applied, err) + } + if applied, err := codes.CompareAndDelete(ctx, "scoped-cas", snapshot.Revision); err != nil || applied { + t.Fatalf("scoped delete applied=%v err=%v", applied, err) + } + if _, found, _ := codes.Get(ctx, "scoped-cas"); !found { + t.Fatal("generic CAS mutated scoped code") + } +} + +func TestCodeStoreRevisionCASPreventsABAAndHasSingleWinner(t *testing.T) { + ctx := context.Background() + codes := NewCodeStore() + record := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016204", + Code: "123456", + Channel: "email_change", + } + if err := codes.Set(ctx, "aba", record, time.Minute); err != nil { + t.Fatal(err) + } + old, found, err := codes.GetSnapshot(ctx, "aba") + if err != nil || !found { + t.Fatalf("old snapshot found=%v err=%v", found, err) + } + if err := codes.Set(ctx, "aba", record, time.Minute); err != nil { + t.Fatal(err) + } + current, found, err := codes.GetSnapshot(ctx, "aba") + if err != nil || !found || current.Revision == old.Revision { + t.Fatalf("replacement snapshot=%+v old=%+v found=%v err=%v", current, old, found, err) + } + if applied, err := codes.CompareAndDelete(ctx, "aba", old.Revision); err != nil || applied { + t.Fatalf("ABA stale delete applied=%v err=%v", applied, err) + } + + const workers = 64 + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + next := current.Record + next.Code = fmt.Sprintf("%06d", index) + applied, err := codes.CompareAndUpdate(ctx, "aba", current.Revision, next) + if err != nil { + errs <- err + return + } + results <- applied + }(i) + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("concurrent CAS update: %v", err) + } + winners := 0 + for applied := range results { + if applied { + winners++ + } + } + if winners != 1 { + t.Fatalf("concurrent CAS update winners=%d, want 1", winners) + } + winner, found, err := codes.GetSnapshot(ctx, "aba") + if err != nil || !found || winner.Revision == current.Revision { + t.Fatalf("winner snapshot=%+v found=%v err=%v", winner, found, err) + } + + results = make(chan bool, workers) + errs = make(chan error, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + applied, err := codes.CompareAndDelete(ctx, "aba", winner.Revision) + if err != nil { + errs <- err + return + } + results <- applied + }() + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("concurrent CAS delete: %v", err) + } + winners = 0 + for applied := range results { + if applied { + winners++ + } + } + if winners != 1 { + t.Fatalf("concurrent CAS delete winners=%d, want 1", winners) + } +} diff --git a/internal/store/memory/code_test.go b/internal/store/memory/code_test.go index b01227b8..30494356 100644 --- a/internal/store/memory/code_test.go +++ b/internal/store/memory/code_test.go @@ -13,6 +13,7 @@ func TestCodeStoreScopedRotationAndSingleConsume(t *testing.T) { ctx := context.Background() codes := NewCodeStore() rec := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, Phone: "15550015001", Code: "12345", Purpose: store.PhoneCodePurposeChangePhone, @@ -68,7 +69,7 @@ func TestCodeStoreScopedRotationAndSingleConsume(t *testing.T) { func TestCodeStoreScopedIsolation(t *testing.T) { ctx := context.Background() codes := NewCodeStore() - a := store.PhoneCode{Phone: "15550015002", Code: "12345", Purpose: store.PhoneCodePurposeChangePhone, UserID: 42, AuthKeyID: [8]byte{1}} + a := store.PhoneCode{Version: store.PhoneCodeVersionCurrent, Phone: "15550015002", Code: "12345", Purpose: store.PhoneCodePurposeChangePhone, UserID: 42, AuthKeyID: [8]byte{1}} b := a b.AuthKeyID = [8]byte{2} if err := codes.Set(ctx, "hash-a", a, time.Minute); err != nil { @@ -90,3 +91,24 @@ func TestCodeStoreScopedIsolation(t *testing.T) { t.Fatal("other scope was removed") } } + +func TestCodeStoreConsumeScopedRejectsAndDeletesLegacyVersion(t *testing.T) { + ctx := context.Background() + codes := NewCodeStore() + rec := store.PhoneCode{ + Version: 0, Phone: "15550015003", Code: "12345", + Purpose: store.PhoneCodePurposeChangePhone, UserID: 43, AuthKeyID: [8]byte{3}, + } + if err := codes.Set(ctx, "legacy-scope", rec, time.Minute); err != nil { + t.Fatal(err) + } + if _, found, err := codes.ConsumeScoped(ctx, "legacy-scope", rec.Scope()); err != nil || found { + t.Fatalf("legacy scoped consume found=%v err=%v, want false/nil", found, err) + } + if _, found, _ := codes.Get(ctx, "legacy-scope"); found { + t.Fatal("legacy scoped code remains after fail-closed consume") + } + if _, ok := codes.scopes[rec.Scope()]; ok { + t.Fatal("legacy scoped index remains after fail-closed consume") + } +} diff --git a/internal/store/memory/login_code.go b/internal/store/memory/login_code.go new file mode 100644 index 00000000..cbef2f63 --- /dev/null +++ b/internal/store/memory/login_code.go @@ -0,0 +1,211 @@ +package memory + +import ( + "context" + "crypto/subtle" + "time" + + "telesrv/internal/store" +) + +func (s *CodeStore) VerifyLogin(_ context.Context, hash, phone, code string, keepForSignUp bool, defaultMaxAttempts int) (store.LoginCodeVerifyResult, error) { + s.mu.Lock() + defer s.mu.Unlock() + + entry, ok := s.liveCodeLocked(hash) + if !ok { + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyMissing}, nil + } + record := entry.code + if record.Version != store.PhoneCodeVersionCurrent { + s.deleteCodeLocked(hash, record) + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyMissing}, nil + } + // Sign-up verification is a terminal state for VerifyLogin. Keep the marker + // for the one caller that already received signUpRequired, but do not report + // a second Accepted result or let later wrong-code calls exhaust it. + if record.SignUpVerified { + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyMissing}, nil + } + if record.Purpose != "" || record.Phone != phone || !loginCodeVerifiable(record) || record.Code == "" || code == "" { + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyInvalid, Record: record}, nil + } + if subtle.ConstantTimeCompare([]byte(record.Code), []byte(code)) != 1 { + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return store.LoginCodeVerifyResult{}, err + } + record.Attempts++ + record.Revision = revision + entry.code = record + maxAttempts := record.MaxAttempts + if maxAttempts <= 0 { + maxAttempts = defaultMaxAttempts + } + if maxAttempts <= 0 { + maxAttempts = 1 + } + if record.Attempts >= maxAttempts { + s.deleteCodeLocked(hash, record) + } else { + s.m[hash] = entry + } + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyInvalid, Record: record}, nil + } + + if keepForSignUp { + if record.IssuedUserID != 0 { + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyInvalid, Record: record}, nil + } + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return store.LoginCodeVerifyResult{}, err + } + record.SignUpVerified = true + record.Revision = revision + entry.code = record + s.m[hash] = entry + } else { + s.deleteCodeLocked(hash, record) + } + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyAccepted, Record: record}, nil +} + +func (s *CodeStore) VerifyScoped(_ context.Context, hash string, scope store.PhoneCodeScope, code string, defaultMaxAttempts int) (store.LoginCodeVerifyResult, error) { + s.mu.Lock() + defer s.mu.Unlock() + + if !scope.Valid() || s.scopes[scope] != hash { + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyMissing}, nil + } + entry, ok := s.liveCodeLocked(hash) + if !ok { + // liveCodeLocked removes the index when it can decode the stored scope; + // also close the stale-index-only case. + if s.scopes[scope] == hash { + delete(s.scopes, scope) + } + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyMissing}, nil + } + record := entry.code + if record.Version != store.PhoneCodeVersionCurrent || record.Scope() != scope || record.SignUpVerified || record.Code == "" { + // Fail closed on a legacy or internally inconsistent record. Clean both + // the scope encoded in the record and the scope that selected this hash. + s.deleteCodeLocked(hash, record) + if s.scopes[scope] == hash { + delete(s.scopes, scope) + } + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyMissing}, nil + } + if code == "" { + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyInvalid, Record: record}, nil + } + if subtle.ConstantTimeCompare([]byte(record.Code), []byte(code)) != 1 { + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return store.LoginCodeVerifyResult{}, err + } + record.Attempts++ + record.Revision = revision + entry.code = record + maxAttempts := record.MaxAttempts + if maxAttempts <= 0 { + maxAttempts = defaultMaxAttempts + } + if maxAttempts <= 0 { + maxAttempts = 1 + } + if record.Attempts >= maxAttempts { + s.deleteCodeLocked(hash, record) + } else { + s.m[hash] = entry + } + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyInvalid, Record: record}, nil + } + + s.deleteCodeLocked(hash, record) + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyAccepted, Record: record}, nil +} + +func (s *CodeStore) ConsumeSignUpVerified(_ context.Context, hash, phone string) (store.PhoneCode, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + + entry, ok := s.liveCodeLocked(hash) + if !ok { + return store.PhoneCode{}, false, nil + } + record := entry.code + if record.Version != store.PhoneCodeVersionCurrent { + s.deleteCodeLocked(hash, record) + return store.PhoneCode{}, false, nil + } + if record.Purpose != "" || record.Phone != phone || !loginCodeVerifiable(record) || record.IssuedUserID != 0 || !record.SignUpVerified { + return store.PhoneCode{}, false, nil + } + s.deleteCodeLocked(hash, record) + return record, true, nil +} + +func (s *CodeStore) TakeLoginCode(_ context.Context, hash, phone string) (store.PhoneCode, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + + entry, ok := s.liveCodeLocked(hash) + if !ok { + return store.PhoneCode{}, false, nil + } + record := entry.code + if record.Version != store.PhoneCodeVersionCurrent { + s.deleteCodeLocked(hash, record) + return store.PhoneCode{}, false, nil + } + if record.SignUpVerified { + return store.PhoneCode{}, false, nil + } + if record.Purpose != "" || record.Phone != phone || !loginCodeTakeable(record) { + return store.PhoneCode{}, false, nil + } + s.deleteCodeLocked(hash, record) + return record, true, nil +} + +func (s *CodeStore) InvalidateLoginCode(_ context.Context, hash, phone string) (bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + + entry, ok := s.liveCodeLocked(hash) + if !ok { + return false, nil + } + record := entry.code + if record.Version != store.PhoneCodeVersionCurrent { + s.deleteCodeLocked(hash, record) + return false, nil + } + if record.Purpose != "" || record.Phone != phone || !loginCodeTakeable(record) { + return false, nil + } + s.deleteCodeLocked(hash, record) + return true, nil +} + +func loginCodeVerifiable(record store.PhoneCode) bool { + return record.Channel == store.PhoneCodeChannelPhone || record.Channel == store.PhoneCodeChannelEmailLogin +} + +func loginCodeTakeable(record store.PhoneCode) bool { + return loginCodeVerifiable(record) || record.Channel == store.PhoneCodeChannelEmailSetupRequired +} + +func (s *CodeStore) liveCodeLocked(hash string) (codeEntry, bool) { + entry, ok := s.m[hash] + if !ok { + return codeEntry{}, false + } + if time.Now().After(entry.expires) { + s.deleteCodeLocked(hash, entry.code) + return codeEntry{}, false + } + return entry, true +} diff --git a/internal/store/memory/login_code_delivery.go b/internal/store/memory/login_code_delivery.go new file mode 100644 index 00000000..6fd06f92 --- /dev/null +++ b/internal/store/memory/login_code_delivery.go @@ -0,0 +1,129 @@ +package memory + +import ( + "context" + "fmt" + "time" + + "telesrv/internal/domain" + "telesrv/internal/store" +) + +type loginCodeDeliveryRecord struct { + userID int64 + codeFingerprint [32]byte + privateMessageID int64 + messageBoxID int + pts int + messageDate int +} + +// LoginCodeDeliveryStore composes the message projection with the same +// durable update-event store consumed by updates.getDifference. Keeping this +// as an explicit dependency prevents tests (and future in-memory runtimes) +// from accidentally creating a visible message without its pts event. +type LoginCodeDeliveryStore struct { + messages *MessageStore + events *UpdateEventStore +} + +func NewLoginCodeDeliveryStore(messages *MessageStore, events *UpdateEventStore) *LoginCodeDeliveryStore { + return &LoginCodeDeliveryStore{messages: messages, events: events} +} + +// DeliverLoginCodeMessage is the in-memory equivalent of the PostgreSQL +// transaction. Message, dialog, pts, durable event and immutable receipt are +// published under the two backing stores' locks; a repeated phone_code_hash +// returns the first snapshot. +func (s *LoginCodeDeliveryStore) DeliverLoginCodeMessage(_ context.Context, req domain.LoginCodeDeliveryRequest) (domain.LoginCodeDeliveryResult, error) { + if s == nil || s.messages == nil || s.events == nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("memory login code delivery: %w: message and update stores are required", domain.ErrLoginCodeDeliveryInvalid) + } + deliveryKey, err := store.LoginCodeDeliveryKey(req.PhoneCodeHash) + if err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + codeFingerprint, err := store.LoginCodeFingerprint(req.PhoneCodeHash, req.Code) + if err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + if req.Date == 0 { + req.Date = int(time.Now().Unix()) + } + if req.ExpiresAt <= int64(req.Date) { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("memory login code receipt expiry: %w: date=%d expires_at=%d", domain.ErrLoginCodeDeliveryInvalid, req.Date, req.ExpiresAt) + } + base, err := domain.OfficialLoginCodeMessage(req.UserID, req.Code, req.Date) + if err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + + // Lock ordering is local and fixed: MessageStore -> UpdateEventStore -> + // DialogStore. No other memory operation takes the first two together. + s.messages.mu.Lock() + defer s.messages.mu.Unlock() + s.events.mu.Lock() + defer s.events.mu.Unlock() + if receipt, ok := s.messages.loginCodeDeliveries[deliveryKey]; ok { + if receipt.userID != req.UserID || !store.SameLoginCodeFingerprint(receipt.codeFingerprint[:], codeFingerprint) { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("memory login code delivery: %w", domain.ErrLoginCodeDeliveryConflict) + } + msg, err := store.RestoreLoginCodeDeliveryMessage( + receipt.userID, + req.Code, + receipt.messageDate, + receipt.privateMessageID, + receipt.messageBoxID, + receipt.pts, + ) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("memory login code delivery replay: %w", err) + } + return domain.LoginCodeDeliveryResult{Message: msg, Created: false}, nil + } + + currentEventPts := 0 + for _, event := range s.events.events[req.UserID] { + if event.Pts > currentEventPts { + currentEventPts = event.Pts + } + } + if messagePts := s.messages.nextPts[req.UserID]; messagePts > currentEventPts { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("memory login code delivery: %w: message pts %d exceeds durable event pts %d", domain.ErrLoginCodeDeliveryInvalid, messagePts, currentEventPts) + } + + base.ID = s.messages.nextBoxIDLocked(req.UserID) + base.UID = s.messages.nextUID + s.messages.nextUID++ + base.Pts = currentEventPts + 1 + + s.messages.nextPts[req.UserID] = base.Pts + s.messages.m[req.UserID] = append(s.messages.m[req.UserID], cloneMessage(base)) + if s.messages.dialogs != nil { + s.messages.dialogs.mu.Lock() + list := s.messages.dialogs.m[req.UserID] + list = upsertMemoryDialog(list, domain.Dialog{ + Peer: base.Peer, + TopMessage: base.ID, + TopMessageDate: base.Date, + UnreadCount: s.messages.privateUnreadCountLocked(req.UserID, base.Peer), + }) + if !hasUser(list.Users, domain.OfficialSystemUserID) { + list.Users = append(list.Users, domain.OfficialSystemUser()) + } + list.Messages = append(list.Messages, cloneMessage(base)) + s.messages.dialogs.m[req.UserID] = list + s.messages.dialogs.mu.Unlock() + } + event := newMessageEvent(base) + s.events.events[req.UserID] = append(s.events.events[req.UserID], cloneUpdateEvent(event)) + s.messages.loginCodeDeliveries[deliveryKey] = loginCodeDeliveryRecord{ + userID: req.UserID, + codeFingerprint: codeFingerprint, + privateMessageID: base.UID, + messageBoxID: base.ID, + pts: base.Pts, + messageDate: base.Date, + } + return domain.LoginCodeDeliveryResult{Message: cloneMessage(base), Created: true}, nil +} diff --git a/internal/store/memory/login_code_delivery_test.go b/internal/store/memory/login_code_delivery_test.go new file mode 100644 index 00000000..8dbf21d9 --- /dev/null +++ b/internal/store/memory/login_code_delivery_test.go @@ -0,0 +1,165 @@ +package memory + +import ( + "context" + "errors" + "reflect" + "sync" + "sync/atomic" + "testing" + + "telesrv/internal/domain" +) + +func TestLoginCodeDeliveryStoreCommitsMessageEventDialogAndReplay(t *testing.T) { + ctx := context.Background() + const userID int64 = 1000000001 + dialogs := NewDialogStore() + messages := NewMessageStore(dialogs) + events := NewUpdateEventStore() + deliveries := NewLoginCodeDeliveryStore(messages, events) + req := domain.LoginCodeDeliveryRequest{ + UserID: userID, + PhoneCodeHash: "phone-code-hash-one", + Code: "12345", + Date: 1700000000, + ExpiresAt: 1700000300, + } + + first, err := deliveries.DeliverLoginCodeMessage(ctx, req) + if err != nil { + t.Fatalf("DeliverLoginCodeMessage: %v", err) + } + if !first.Created || first.Message.ID != 1 || first.Message.UID != 1 || first.Message.Pts != 1 || first.Message.Out || + first.Message.OwnerUserID != userID || first.Message.Peer.ID != domain.OfficialSystemUserID || first.Message.From.ID != domain.OfficialSystemUserID { + t.Fatalf("first delivery = %+v, want first incoming 777000 message", first) + } + if len(messages.m[userID]) != 1 || !reflect.DeepEqual(messages.m[userID][0], first.Message) { + t.Fatalf("message projection = %+v, want committed message", messages.m[userID]) + } + if len(events.events[userID]) != 1 { + t.Fatalf("durable events = %+v, want one new_message", events.events[userID]) + } + event := events.events[userID][0] + if event.Type != domain.UpdateEventNewMessage || event.Pts != first.Message.Pts || event.PtsCount != 1 || !reflect.DeepEqual(event.Message, first.Message) { + t.Fatalf("event = %+v, want message-identical new_message", event) + } + list := dialogs.m[userID] + if len(list.Dialogs) != 1 || list.Dialogs[0].Peer.ID != domain.OfficialSystemUserID || list.Dialogs[0].TopMessage != first.Message.ID || list.Dialogs[0].UnreadCount != 1 { + t.Fatalf("dialog projection = %+v, want unread 777000 dialog", list.Dialogs) + } + if len(list.Users) != 1 || list.Users[0].ID != domain.OfficialSystemUserID { + t.Fatalf("dialog users = %+v, want official system user", list.Users) + } + + replayReq := req + replayReq.Date++ + replay, err := deliveries.DeliverLoginCodeMessage(ctx, replayReq) + if err != nil { + t.Fatalf("replay DeliverLoginCodeMessage: %v", err) + } + if replay.Created || !reflect.DeepEqual(replay.Message, first.Message) { + t.Fatalf("replay = %+v, want immutable first result %+v", replay, first) + } + if len(messages.m[userID]) != 1 || len(events.events[userID]) != 1 || len(messages.loginCodeDeliveries) != 1 { + t.Fatalf("replay created facts: messages=%d events=%d receipts=%d", len(messages.m[userID]), len(events.events[userID]), len(messages.loginCodeDeliveries)) + } + + second, err := deliveries.DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{ + UserID: userID, + PhoneCodeHash: "phone-code-hash-two", + Code: "67890", + Date: 1700000010, + ExpiresAt: 1700000310, + }) + if err != nil { + t.Fatalf("second distinct delivery: %v", err) + } + if !second.Created || second.Message.ID != 2 || second.Message.UID != 2 || second.Message.Pts != 2 || len(events.events[userID]) != 2 { + t.Fatalf("second delivery = %+v events=%+v, want contiguous allocations", second, events.events[userID]) + } + if got := dialogs.m[userID].Dialogs[0].UnreadCount; got != 2 { + t.Fatalf("dialog unread = %d, want 2", got) + } +} + +func TestLoginCodeDeliveryStoreConcurrentReplayAndConflict(t *testing.T) { + ctx := context.Background() + const userID int64 = 1000000002 + messages := NewMessageStore(NewDialogStore()) + events := NewUpdateEventStore() + deliveries := NewLoginCodeDeliveryStore(messages, events) + req := domain.LoginCodeDeliveryRequest{ + UserID: userID, + PhoneCodeHash: "concurrent-phone-code-hash", + Code: "24680", + Date: 1700000100, + ExpiresAt: 1700000400, + } + + const workers = 32 + var created atomic.Int32 + results := make(chan domain.LoginCodeDeliveryResult, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + got, err := deliveries.DeliverLoginCodeMessage(ctx, req) + if err != nil { + errs <- err + return + } + if got.Created { + created.Add(1) + } + results <- got + }() + } + wg.Wait() + close(errs) + close(results) + for err := range errs { + t.Fatalf("concurrent delivery: %v", err) + } + if created.Load() != 1 { + t.Fatalf("created calls = %d, want exactly 1", created.Load()) + } + for got := range results { + if got.Message.ID != 1 || got.Message.UID != 1 || got.Message.Pts != 1 { + t.Fatalf("concurrent result = %+v, want the same first allocation", got) + } + } + if len(messages.m[userID]) != 1 || len(events.events[userID]) != 1 || len(messages.loginCodeDeliveries) != 1 { + t.Fatalf("concurrent facts: messages=%d events=%d receipts=%d", len(messages.m[userID]), len(events.events[userID]), len(messages.loginCodeDeliveries)) + } + + changedCode := req + changedCode.Code = "13579" + if _, err := deliveries.DeliverLoginCodeMessage(ctx, changedCode); !errors.Is(err, domain.ErrLoginCodeDeliveryConflict) { + t.Fatalf("changed-code replay err = %v, want ErrLoginCodeDeliveryConflict", err) + } + changedUser := req + changedUser.UserID++ + if _, err := deliveries.DeliverLoginCodeMessage(ctx, changedUser); !errors.Is(err, domain.ErrLoginCodeDeliveryConflict) { + t.Fatalf("changed-user replay err = %v, want ErrLoginCodeDeliveryConflict", err) + } + if len(messages.m[userID]) != 1 || len(events.events[userID]) != 1 { + t.Fatal("conflicting replay changed committed facts") + } +} + +func TestLoginCodeDeliveryStoreRequiresSharedEventStore(t *testing.T) { + messages := NewMessageStore() + _, err := NewLoginCodeDeliveryStore(messages, nil).DeliverLoginCodeMessage(context.Background(), domain.LoginCodeDeliveryRequest{ + UserID: 1000000003, + PhoneCodeHash: "missing-event-store", + Code: "12345", + Date: 1700000200, + ExpiresAt: 1700000500, + }) + if !errors.Is(err, domain.ErrLoginCodeDeliveryInvalid) { + t.Fatalf("missing event store err = %v, want ErrLoginCodeDeliveryInvalid", err) + } +} diff --git a/internal/store/memory/login_code_invalidate_test.go b/internal/store/memory/login_code_invalidate_test.go new file mode 100644 index 00000000..74904ff2 --- /dev/null +++ b/internal/store/memory/login_code_invalidate_test.go @@ -0,0 +1,106 @@ +package memory + +import ( + "context" + "sync" + "testing" + "time" + + "telesrv/internal/store" +) + +func TestCodeStoreAtomicLoginInvalidation(t *testing.T) { + ctx := context.Background() + const phone = "15550016011" + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: phone, + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + MaxAttempts: 5, + } + } + + t.Run("owner cleanup may delete a terminal sign-up marker", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "invalidate-marker", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + verified, err := codes.VerifyLogin(ctx, "invalidate-marker", phone, "12345", true, 5) + if err != nil || verified.Status != store.LoginCodeVerifyAccepted || !verified.Record.SignUpVerified { + t.Fatalf("mark sign-up = %+v err=%v", verified, err) + } + if removed, err := codes.InvalidateLoginCode(ctx, "invalidate-marker", "15550016999"); err != nil || removed { + t.Fatalf("cross-phone invalidate removed=%v err=%v", removed, err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, "invalidate-marker", "15550016999"); err != nil || found { + t.Fatalf("cross-phone consume found=%v err=%v", found, err) + } + if removed, err := codes.InvalidateLoginCode(ctx, "invalidate-marker", phone); err != nil || !removed { + t.Fatalf("owner invalidate removed=%v err=%v", removed, err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, "invalidate-marker", phone); err != nil || found { + t.Fatalf("consume after invalidate found=%v err=%v", found, err) + } + }) + + t.Run("legacy records fail closed", func(t *testing.T) { + codes := NewCodeStore() + legacy := newRecord() + legacy.Version = 0 + if err := codes.Set(ctx, "invalidate-legacy", legacy, time.Minute); err != nil { + t.Fatal(err) + } + if removed, err := codes.InvalidateLoginCode(ctx, "invalidate-legacy", phone); err != nil || removed { + t.Fatalf("legacy invalidate removed=%v err=%v, want false", removed, err) + } + if _, found, err := codes.Get(ctx, "invalidate-legacy"); err != nil || found { + t.Fatalf("legacy record found=%v err=%v after fail-closed invalidate", found, err) + } + }) + + t.Run("invalidate and sign-up consume have one winner", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "invalidate-race", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + if verified, err := codes.VerifyLogin(ctx, "invalidate-race", phone, "12345", true, 5); err != nil || verified.Status != store.LoginCodeVerifyAccepted { + t.Fatalf("mark sign-up = %+v err=%v", verified, err) + } + + const workers = 64 + results := make(chan bool, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(invalidate bool) { + defer wg.Done() + if invalidate { + removed, err := codes.InvalidateLoginCode(ctx, "invalidate-race", phone) + if err != nil { + t.Errorf("InvalidateLoginCode: %v", err) + } + results <- removed + return + } + _, found, err := codes.ConsumeSignUpVerified(ctx, "invalidate-race", phone) + if err != nil { + t.Errorf("ConsumeSignUpVerified: %v", err) + } + results <- found + }(i%2 == 0) + } + wg.Wait() + close(results) + winners := 0 + for won := range results { + if won { + winners++ + } + } + if winners != 1 { + t.Fatalf("invalidate/consume winners=%d, want 1", winners) + } + }) +} diff --git a/internal/store/memory/login_code_state_test.go b/internal/store/memory/login_code_state_test.go new file mode 100644 index 00000000..fc02dfa2 --- /dev/null +++ b/internal/store/memory/login_code_state_test.go @@ -0,0 +1,447 @@ +package memory + +import ( + "context" + "sync" + "testing" + "time" + + "telesrv/internal/store" +) + +func TestCodeStoreAtomicLoginStateMachine(t *testing.T) { + ctx := context.Background() + const phone = "15550016001" + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + IssuedUserID: 1000000001, + Phone: phone, + Code: "12345", + Channel: "phone", + MaxAttempts: 2, + } + } + + t.Run("version mismatch fails closed for every atomic entry", func(t *testing.T) { + codes := NewCodeStore() + legacy := newRecord() + legacy.Version = 0 + if err := codes.Set(ctx, "legacy-verify", legacy, time.Minute); err != nil { + t.Fatal(err) + } + result, err := codes.VerifyLogin(ctx, "legacy-verify", phone, legacy.Code, false, 5) + if err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("legacy VerifyLogin = %+v err=%v, want Missing", result, err) + } + if _, found, _ := codes.Get(ctx, "legacy-verify"); found { + t.Fatal("legacy VerifyLogin record was not deleted") + } + + unknown := newRecord() + unknown.Version = store.PhoneCodeVersionCurrent + 1 + if err := codes.Set(ctx, "unknown-take", unknown, time.Minute); err != nil { + t.Fatal(err) + } + if _, found, err := codes.TakeLoginCode(ctx, "unknown-take", phone); err != nil || found { + t.Fatalf("unknown TakeLoginCode found=%v err=%v, want false", found, err) + } + if _, found, _ := codes.Get(ctx, "unknown-take"); found { + t.Fatal("unknown TakeLoginCode record was not deleted") + } + + legacy.SignUpVerified = true + if err := codes.Set(ctx, "legacy-signup", legacy, time.Minute); err != nil { + t.Fatal(err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, "legacy-signup", phone); err != nil || found { + t.Fatalf("legacy ConsumeSignUpVerified found=%v err=%v, want false", found, err) + } + if _, found, _ := codes.Get(ctx, "legacy-signup"); found { + t.Fatal("legacy sign-up marker was not deleted") + } + }) + + t.Run("scope mismatch does not burn victim attempts", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord() + if err := codes.Set(ctx, "scope", record, time.Minute); err != nil { + t.Fatal(err) + } + result, err := codes.VerifyLogin(ctx, "scope", "15550016999", record.Code, false, 5) + if err != nil || result.Status != store.LoginCodeVerifyInvalid || result.Record.Attempts != 0 { + t.Fatalf("wrong-phone VerifyLogin = %+v err=%v", result, err) + } + stored, found, err := codes.Get(ctx, "scope") + if err != nil || !found || stored.Attempts != 0 { + t.Fatalf("wrong-phone stored=%+v found=%v err=%v", stored, found, err) + } + if _, found, err := codes.TakeLoginCode(ctx, "scope", "15550016999"); err != nil || found { + t.Fatalf("cross-phone TakeLoginCode found=%v err=%v", found, err) + } + + scoped := newRecord() + scoped.Purpose = store.PhoneCodePurposeChangePhone + scoped.UserID = 42 + scoped.AuthKeyID = [8]byte{1} + if err := codes.Set(ctx, "scoped", scoped, time.Minute); err != nil { + t.Fatal(err) + } + result, err = codes.VerifyLogin(ctx, "scoped", phone, scoped.Code, false, 5) + if err != nil || result.Status != store.LoginCodeVerifyInvalid { + t.Fatalf("scoped VerifyLogin = %+v err=%v, want Invalid", result, err) + } + if _, found, err := codes.TakeLoginCode(ctx, "scoped", phone); err != nil || found { + t.Fatalf("scoped TakeLoginCode found=%v err=%v", found, err) + } + if _, found, _ := codes.Get(ctx, "scoped"); !found { + t.Fatal("login operations deleted a scoped change-phone code") + } + }) + + t.Run("wrong code increments atomically and threshold deletes", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord() + if err := codes.Set(ctx, "wrong", record, time.Minute); err != nil { + t.Fatal(err) + } + first, err := codes.VerifyLogin(ctx, "wrong", phone, "00000", false, 9) + if err != nil || first.Status != store.LoginCodeVerifyInvalid || first.Record.Attempts != 1 { + t.Fatalf("first wrong code = %+v err=%v", first, err) + } + stored, found, err := codes.Get(ctx, "wrong") + if err != nil || !found || stored.Attempts != 1 { + t.Fatalf("stored after first wrong = %+v found=%v err=%v", stored, found, err) + } + second, err := codes.VerifyLogin(ctx, "wrong", phone, "00000", false, 9) + if err != nil || second.Status != store.LoginCodeVerifyInvalid || second.Record.Attempts != 2 { + t.Fatalf("threshold wrong code = %+v err=%v", second, err) + } + if _, found, _ := codes.Get(ctx, "wrong"); found { + t.Fatal("threshold-exhausted code remains") + } + after, err := codes.VerifyLogin(ctx, "wrong", phone, record.Code, false, 9) + if err != nil || after.Status != store.LoginCodeVerifyMissing { + t.Fatalf("verify after exhaustion = %+v err=%v, want Missing", after, err) + } + + fallback := newRecord() + fallback.MaxAttempts = 0 + if err := codes.Set(ctx, "fallback", fallback, time.Minute); err != nil { + t.Fatal(err) + } + if got, err := codes.VerifyLogin(ctx, "fallback", phone, "bad", false, 1); err != nil || got.Status != store.LoginCodeVerifyInvalid { + t.Fatalf("default threshold verify = %+v err=%v", got, err) + } + if _, found, _ := codes.Get(ctx, "fallback"); found { + t.Fatal("default threshold did not delete code") + } + }) + + t.Run("accepted consume and sign-up marker are terminal", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord() + if err := codes.Set(ctx, "consume", record, time.Minute); err != nil { + t.Fatal(err) + } + accepted, err := codes.VerifyLogin(ctx, "consume", phone, record.Code, false, 5) + if err != nil || accepted.Status != store.LoginCodeVerifyAccepted || accepted.Record.SignUpVerified { + t.Fatalf("consume verify = %+v err=%v", accepted, err) + } + if _, found, _ := codes.Get(ctx, "consume"); found { + t.Fatal("accepted existing-user code remains") + } + + if err := codes.Set(ctx, "issued-existing", record, time.Minute); err != nil { + t.Fatal(err) + } + wrongScope, err := codes.VerifyLogin(ctx, "issued-existing", phone, record.Code, true, 5) + if err != nil || wrongScope.Status != store.LoginCodeVerifyInvalid { + t.Fatalf("existing-issued keep-for-signup = %+v err=%v, want Invalid", wrongScope, err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, "issued-existing", phone); err != nil || found { + t.Fatalf("existing-issued sign-up consume found=%v err=%v", found, err) + } + + signUpRecord := record + signUpRecord.IssuedUserID = 0 + if err := codes.Set(ctx, "signup", signUpRecord, time.Minute); err != nil { + t.Fatal(err) + } + expires := codes.m["signup"].expires + marked, err := codes.VerifyLogin(ctx, "signup", phone, record.Code, true, 5) + if err != nil || marked.Status != store.LoginCodeVerifyAccepted || !marked.Record.SignUpVerified { + t.Fatalf("sign-up verify = %+v err=%v", marked, err) + } + if got := codes.m["signup"]; !got.code.SignUpVerified || !got.expires.Equal(expires) { + t.Fatalf("sign-up marker=%+v expiry=%v, want marker with unchanged %v", got.code, got.expires, expires) + } + repeated, err := codes.VerifyLogin(ctx, "signup", phone, record.Code, true, 5) + if err != nil || repeated.Status != store.LoginCodeVerifyMissing { + t.Fatalf("repeated sign-up verify = %+v err=%v, want terminal Missing", repeated, err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, "signup", "15550016999"); err != nil || found { + t.Fatalf("cross-phone sign-up consume found=%v err=%v", found, err) + } + if _, found, err := codes.TakeLoginCode(ctx, "signup", phone); err != nil || found { + t.Fatalf("terminal marker take found=%v err=%v, want false", found, err) + } + consumed, found, err := codes.ConsumeSignUpVerified(ctx, "signup", phone) + if err != nil || !found || !consumed.SignUpVerified || consumed.Code != signUpRecord.Code || consumed.IssuedUserID != 0 { + t.Fatalf("sign-up consume = %+v found=%v err=%v", consumed, found, err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, "signup", phone); err != nil || found { + t.Fatalf("second sign-up consume found=%v err=%v", found, err) + } + }) + + t.Run("take returns the removed record exactly once", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord() + if err := codes.Set(ctx, "take", record, time.Minute); err != nil { + t.Fatal(err) + } + expected, found, err := codes.Get(ctx, "take") + if err != nil || !found { + t.Fatalf("load take record found=%v err=%v", found, err) + } + if _, found, err := codes.TakeLoginCode(ctx, "take", "15550016999"); err != nil || found { + t.Fatalf("cross-phone take found=%v err=%v", found, err) + } + taken, found, err := codes.TakeLoginCode(ctx, "take", phone) + if err != nil || !found || taken != expected { + t.Fatalf("take = %+v found=%v err=%v, want %+v", taken, found, err, expected) + } + if _, found, err := codes.TakeLoginCode(ctx, "take", phone); err != nil || found { + t.Fatalf("second take found=%v err=%v", found, err) + } + }) +} + +func TestCodeStoreAtomicLoginConcurrency(t *testing.T) { + ctx := context.Background() + const ( + phone = "15550016002" + workers = 64 + ) + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: phone, + Code: "12345", + Channel: "phone", + MaxAttempts: 7, + } + } + + t.Run("consume verify has one accepted", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "verify-race", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentMemoryVerify(t, codes, "verify-race", phone, "12345", false, workers) + if statuses[store.LoginCodeVerifyAccepted] != 1 || statuses[store.LoginCodeVerifyMissing] != workers-1 || statuses[store.LoginCodeVerifyInvalid] != 0 { + t.Fatalf("verify race statuses = %+v", statuses) + } + }) + + t.Run("mark and consume each have one winner", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "signup-race", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentMemoryVerify(t, codes, "signup-race", phone, "12345", true, workers) + if statuses[store.LoginCodeVerifyAccepted] != 1 || statuses[store.LoginCodeVerifyMissing] != workers-1 { + t.Fatalf("sign-up verify race statuses = %+v", statuses) + } + found := concurrentMemoryConsumeSignUp(t, codes, "signup-race", phone, workers) + if found != 1 { + t.Fatalf("sign-up consumes = %d, want 1", found) + } + }) + + t.Run("take has one winner", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "take-race", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + found := concurrentMemoryTake(t, codes, "take-race", phone, workers) + if found != 1 { + t.Fatalf("takes = %d, want 1", found) + } + }) + + t.Run("wrong attempts cannot be lost", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "wrong-race", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentMemoryVerify(t, codes, "wrong-race", phone, "00000", false, workers) + if statuses[store.LoginCodeVerifyInvalid] != 7 || statuses[store.LoginCodeVerifyMissing] != workers-7 { + t.Fatalf("wrong-code race statuses = %+v, want 7 Invalid then Missing", statuses) + } + if _, found, _ := codes.Get(ctx, "wrong-race"); found { + t.Fatal("wrong-code race left an exhausted code") + } + }) + + t.Run("verify and cancel-resend take share one winner", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "mixed-race", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + results := make(chan bool, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(take bool) { + defer wg.Done() + if take { + _, found, err := codes.TakeLoginCode(ctx, "mixed-race", phone) + if err != nil { + t.Errorf("TakeLoginCode: %v", err) + } + results <- found + return + } + verified, err := codes.VerifyLogin(ctx, "mixed-race", phone, "12345", false, 5) + if err != nil { + t.Errorf("VerifyLogin: %v", err) + } + results <- verified.Status == store.LoginCodeVerifyAccepted + }(i%2 == 0) + } + wg.Wait() + close(results) + winners := 0 + for won := range results { + if won { + winners++ + } + } + if winners != 1 { + t.Fatalf("mixed verify/take winners = %d, want 1", winners) + } + }) + + t.Run("sign-up mark and take share one winner", func(t *testing.T) { + codes := NewCodeStore() + if err := codes.Set(ctx, "mixed-signup-race", newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + results := make(chan bool, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(take bool) { + defer wg.Done() + if take { + _, found, err := codes.TakeLoginCode(ctx, "mixed-signup-race", phone) + if err != nil { + t.Errorf("TakeLoginCode: %v", err) + } + results <- found + return + } + verified, err := codes.VerifyLogin(ctx, "mixed-signup-race", phone, "12345", true, 5) + if err != nil { + t.Errorf("VerifyLogin: %v", err) + } + results <- verified.Status == store.LoginCodeVerifyAccepted + }(i%2 == 0) + } + wg.Wait() + close(results) + winners := 0 + for won := range results { + if won { + winners++ + } + } + if winners != 1 { + t.Fatalf("mixed sign-up/take winners = %d, want 1", winners) + } + }) +} + +func concurrentMemoryVerify(t *testing.T, codes *CodeStore, hash, phone, code string, keep bool, workers int) map[store.LoginCodeVerifyStatus]int { + t.Helper() + ctx := context.Background() + results := make(chan store.LoginCodeVerifyStatus, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + result, err := codes.VerifyLogin(ctx, hash, phone, code, keep, 5) + if err != nil { + t.Errorf("VerifyLogin: %v", err) + return + } + results <- result.Status + }() + } + wg.Wait() + close(results) + counts := make(map[store.LoginCodeVerifyStatus]int) + for status := range results { + counts[status]++ + } + return counts +} + +func concurrentMemoryTake(t *testing.T, codes *CodeStore, hash, phone string, workers int) int { + t.Helper() + ctx := context.Background() + results := make(chan bool, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, found, err := codes.TakeLoginCode(ctx, hash, phone) + if err != nil { + t.Errorf("TakeLoginCode: %v", err) + return + } + results <- found + }() + } + wg.Wait() + close(results) + foundCount := 0 + for found := range results { + if found { + foundCount++ + } + } + return foundCount +} + +func concurrentMemoryConsumeSignUp(t *testing.T, codes *CodeStore, hash, phone string, workers int) int { + t.Helper() + ctx := context.Background() + results := make(chan bool, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, found, err := codes.ConsumeSignUpVerified(ctx, hash, phone) + if err != nil { + t.Errorf("ConsumeSignUpVerified: %v", err) + return + } + results <- found + }() + } + wg.Wait() + close(results) + foundCount := 0 + for found := range results { + if found { + foundCount++ + } + } + return foundCount +} diff --git a/internal/store/memory/message_delete.go b/internal/store/memory/message_delete.go index 0e4a0980..0bd00140 100644 --- a/internal/store/memory/message_delete.go +++ b/internal/store/memory/message_delete.go @@ -37,9 +37,12 @@ func (s *MessageStore) DeleteMessages(_ context.Context, req domain.DeleteMessag } type deletedMemoryMessage struct { - userID int64 - peer domain.Peer - id int + userID int64 + peer domain.Peer + id int + privateMessageID int64 + messageSenderID int64 + randomID int64 } func (s *MessageStore) finishMemoryDeleteLocked(res domain.DeleteMessagesResult, deleted []deletedMemoryMessage, date int, preserveEmptyDialogs bool) domain.DeleteMessagesResult { @@ -83,6 +86,19 @@ func (s *MessageStore) finishMemoryDeleteLocked(res domain.DeleteMessagesResult, Date: date, MessageIDs: ids, } + for _, row := range deleted { + if row.userID != userID || row.messageSenderID != userID || row.randomID == 0 || row.privateMessageID == 0 { + continue + } + key := privateSendDedupKey{senderUserID: userID, randomID: row.randomID} + record, ok := s.privateSendDedup[key] + if !ok { + continue + } + cloned := cloneUpdateEvent(event) + record.senderDeleteEvent = &cloned + s.privateSendDedup[key] = record + } res.Deleted = append(res.Deleted, domain.DeletedMessagesForUser{ UserID: userID, MessageIDs: ids, diff --git a/internal/store/memory/message_helpers.go b/internal/store/memory/message_helpers.go index 24455764..a8d0b08d 100644 --- a/internal/store/memory/message_helpers.go +++ b/internal/store/memory/message_helpers.go @@ -21,7 +21,14 @@ func (s *MessageStore) deleteMemoryMessagesLocked(userID int64, limit int, match more = true continue } - deleted = append(deleted, deletedMemoryMessage{userID: userID, peer: msg.Peer, id: msg.ID}) + deleted = append(deleted, deletedMemoryMessage{ + userID: userID, + peer: msg.Peer, + id: msg.ID, + privateMessageID: msg.UID, + messageSenderID: msg.From.ID, + randomID: msg.RandomID, + }) if msg.UID != 0 { revokeUIDs[msg.UID] = struct{}{} } @@ -45,7 +52,14 @@ func (s *MessageStore) deleteMemoryMessagesByUIDLocked(uids map[int64]struct{}, kept := messages[:0] for _, msg := range messages { if _, ok := uids[msg.UID]; ok { - deleted = append(deleted, deletedMemoryMessage{userID: userID, peer: msg.Peer, id: msg.ID}) + deleted = append(deleted, deletedMemoryMessage{ + userID: userID, + peer: msg.Peer, + id: msg.ID, + privateMessageID: msg.UID, + messageSenderID: msg.From.ID, + randomID: msg.RandomID, + }) continue } kept = append(kept, msg) diff --git a/internal/store/memory/message_idempotency_test.go b/internal/store/memory/message_idempotency_test.go new file mode 100644 index 00000000..270e626e --- /dev/null +++ b/internal/store/memory/message_idempotency_test.go @@ -0,0 +1,164 @@ +package memory + +import ( + "context" + "errors" + "testing" + + "telesrv/internal/domain" +) + +func TestMessageStorePrivateRandomIDConflictAndReplayFacts(t *testing.T) { + ctx := context.Background() + messages := NewMessageStore() + base := domain.SendPrivateTextRequest{ + SenderUserID: 1001, RecipientUserID: 1002, RandomID: 501, + Message: "immutable", Date: 1700000000, + } + first, err := messages.SendPrivateText(ctx, base) + if err != nil { + t.Fatalf("first send: %v", err) + } + replay := base + replay.Date++ + replay.OriginSessionID = 77 + replay.RecipientBlocked = true + duplicate, err := messages.SendPrivateText(ctx, replay) + if err != nil { + t.Fatalf("exact replay: %v", err) + } + if !duplicate.Duplicate || duplicate.SenderMessage.ID != first.SenderMessage.ID || duplicate.RecipientMessage.ID != first.RecipientMessage.ID { + t.Fatalf("exact replay = %+v, want original delivered boxes", duplicate) + } + + tests := []struct { + name string + mutate func(*domain.SendPrivateTextRequest) + }{ + {name: "peer", mutate: func(req *domain.SendPrivateTextRequest) { req.RecipientUserID = 1003 }}, + {name: "body", mutate: func(req *domain.SendPrivateTextRequest) { req.Message = "changed" }}, + {name: "media", mutate: func(req *domain.SendPrivateTextRequest) { + req.Media = &domain.MessageMedia{ + Kind: domain.MessageMediaKindContact, + Contact: &domain.MessageContact{ + PhoneNumber: "+10000000000", + FirstName: "Changed", + }, + } + }}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + req := base + tc.mutate(&req) + if _, err := messages.SendPrivateText(ctx, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("conflicting replay err = %v, want ErrMessageRandomIDDuplicate", err) + } + }) + } +} + +func TestMessageStorePrivateRandomIDSelfAndBlockedReplay(t *testing.T) { + ctx := context.Background() + messages := NewMessageStore() + selfReq := domain.SendPrivateTextRequest{ + SenderUserID: 2001, RecipientUserID: 2001, RandomID: 601, + Message: "saved note", Date: 1700000100, + } + self, err := messages.SendPrivateText(ctx, selfReq) + if err != nil { + t.Fatalf("self first: %v", err) + } + selfReq.Date++ + selfReplay, err := messages.SendPrivateText(ctx, selfReq) + if err != nil { + t.Fatalf("self replay: %v", err) + } + if !selfReplay.Duplicate || selfReplay.SenderMessage.ID != self.SenderMessage.ID || selfReplay.RecipientMessage.ID != self.SenderMessage.ID { + t.Fatalf("self replay = %+v, want original single box", selfReplay) + } + + blockedReq := domain.SendPrivateTextRequest{ + SenderUserID: 2002, RecipientUserID: 2003, RandomID: 602, + Message: "blocked", Date: 1700000110, RecipientBlocked: true, + } + blocked, err := messages.SendPrivateText(ctx, blockedReq) + if err != nil { + t.Fatalf("blocked first: %v", err) + } + if blocked.RecipientMessage.ID != 0 { + t.Fatalf("blocked recipient = %+v, want empty", blocked.RecipientMessage) + } + blockedReq.Date++ + blockedReq.RecipientBlocked = false + blockedReplay, err := messages.SendPrivateText(ctx, blockedReq) + if err != nil { + t.Fatalf("blocked replay: %v", err) + } + if !blockedReplay.Duplicate || blockedReplay.SenderMessage.ID != blocked.SenderMessage.ID || blockedReplay.RecipientMessage.ID != 0 { + t.Fatalf("blocked replay = %+v, want original sender-only result", blockedReplay) + } +} + +func TestMessageStorePrivateRandomIDReplayUsesCurrentSnapshotAndDurableDeleteMemory(t *testing.T) { + ctx := context.Background() + messages := NewMessageStore() + req := domain.SendPrivateTextRequest{ + SenderUserID: 3001, RecipientUserID: 3002, RandomID: 701, + Message: "original", Date: 1700000200, + } + first, err := messages.SendPrivateText(ctx, req) + if err != nil { + t.Fatalf("first send: %v", err) + } + edited, err := messages.EditMessage(ctx, domain.EditMessageRequest{ + OwnerUserID: req.SenderUserID, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: req.RecipientUserID}, + ID: first.SenderMessage.ID, + Message: "edited projection", + EditDate: 1700000201, + }) + if err != nil { + t.Fatalf("edit message: %v", err) + } + senderPtsBeforeReplay := messages.nextPts[req.SenderUserID] + recipientPtsBeforeReplay := messages.nextPts[req.RecipientUserID] + replay, err := messages.SendPrivateText(ctx, req) + if err != nil { + t.Fatalf("replay after edit: %v", err) + } + if replay.SenderMessage.ID != first.SenderMessage.ID || replay.SenderMessage.Pts != edited.Self().Message.Pts || replay.SenderMessage.Body != "edited projection" { + t.Fatalf("replay after edit = %+v, want current visible snapshot", replay.SenderMessage) + } + if replay.SenderEvent.Pts != first.SenderEvent.Pts || replay.ReplayDeleteEvent != nil { + t.Fatalf("replay after edit event = %+v delete=%+v, want original send pts and no delete", replay.SenderEvent, replay.ReplayDeleteEvent) + } + if messages.nextPts[req.SenderUserID] != senderPtsBeforeReplay || messages.nextPts[req.RecipientUserID] != recipientPtsBeforeReplay { + t.Fatalf("edit replay advanced pts sender/recipient = %d/%d, want %d/%d", messages.nextPts[req.SenderUserID], messages.nextPts[req.RecipientUserID], senderPtsBeforeReplay, recipientPtsBeforeReplay) + } + deleted, err := messages.DeleteMessages(ctx, domain.DeleteMessagesRequest{ + OwnerUserID: req.SenderUserID, + IDs: []int{first.SenderMessage.ID}, + Revoke: true, + Date: 1700000202, + }) + if err != nil { + t.Fatalf("delete message: %v", err) + } + senderPtsBeforeReplay = messages.nextPts[req.SenderUserID] + recipientPtsBeforeReplay = messages.nextPts[req.RecipientUserID] + replay, err = messages.SendPrivateText(ctx, req) + if err != nil { + t.Fatalf("replay after delete: %v", err) + } + if replay.SenderMessage.ID != first.SenderMessage.ID || replay.SenderMessage.Pts != first.SenderMessage.Pts || replay.SenderMessage.Body != "original" { + t.Fatalf("replay after delete = %+v, want immutable first snapshot", replay.SenderMessage) + } + if replay.ReplayDeleteEvent == nil || replay.ReplayDeleteEvent.Pts != deleted.Self().Event.Pts || + len(replay.ReplayDeleteEvent.MessageIDs) != 1 || replay.ReplayDeleteEvent.MessageIDs[0] != first.SenderMessage.ID { + t.Fatalf("replay delete event = %+v, want durable delete %+v", replay.ReplayDeleteEvent, deleted.Self().Event) + } + if messages.nextPts[req.SenderUserID] != senderPtsBeforeReplay || messages.nextPts[req.RecipientUserID] != recipientPtsBeforeReplay { + t.Fatalf("delete replay advanced pts sender/recipient = %d/%d, want %d/%d", messages.nextPts[req.SenderUserID], messages.nextPts[req.RecipientUserID], senderPtsBeforeReplay, recipientPtsBeforeReplay) + } +} diff --git a/internal/store/memory/message_send.go b/internal/store/memory/message_send.go index 1a1e3f3f..336b52cc 100644 --- a/internal/store/memory/message_send.go +++ b/internal/store/memory/message_send.go @@ -2,10 +2,25 @@ package memory import ( "context" + "fmt" "telesrv/internal/domain" + "telesrv/internal/store" "time" ) +type privateSendDedupKey struct { + senderUserID int64 + randomID int64 +} + +type privateSendDedupRecord struct { + recipientUserID int64 + senderSnapshot []byte + recipientMessage domain.Message + fingerprint []byte + senderDeleteEvent *domain.UpdateEvent +} + func (s *MessageStore) Create(_ context.Context, msg domain.Message) (domain.Message, error) { s.mu.Lock() defer s.mu.Unlock() @@ -30,29 +45,19 @@ func (s *MessageStore) Create(_ context.Context, msg domain.Message) (domain.Mes } func (s *MessageStore) SendPrivateText(_ context.Context, req domain.SendPrivateTextRequest) (domain.SendPrivateTextResult, error) { + fingerprint, err := store.PrivateSendFingerprint(req) + if err != nil { + return domain.SendPrivateTextResult{}, err + } s.mu.Lock() defer s.mu.Unlock() - for _, msg := range s.m[req.SenderUserID] { - if msg.RandomID != 0 && msg.RandomID == req.RandomID { - recipient := domain.Message{} - if req.SenderUserID != req.RecipientUserID { - for _, peerMsg := range s.m[req.RecipientUserID] { - if peerMsg.UID == msg.UID { - recipient = peerMsg - break - } - } - } else { - recipient = msg - } - return domain.SendPrivateTextResult{ - SenderMessage: cloneMessage(msg), - RecipientMessage: cloneMessage(recipient), - SenderEvent: newMessageEvent(msg), - RecipientEvent: newMessageEvent(recipient), - Duplicate: true, - }, nil - } + if replay, found, err := s.lookupPrivateSendReplayLocked(domain.PrivateSendReplayRequest{ + SenderUserID: req.SenderUserID, + RecipientUserID: req.RecipientUserID, + RandomID: req.RandomID, + IdempotencyFingerprint: fingerprint, + }); err != nil || found { + return replay, err } if req.Date == 0 { req.Date = int(time.Now().Unix()) @@ -109,10 +114,25 @@ func (s *MessageStore) SendPrivateText(_ context.Context, req domain.SendPrivate recipient.Pts = s.nextPtsLocked(req.RecipientUserID) recipient.MediaUnread = req.Media.HasUnreadPayload() } + var senderSnapshot []byte + if req.RandomID != 0 { + senderSnapshot, err = store.EncodePrivateSendSnapshot(sender) + if err != nil { + return domain.SendPrivateTextResult{}, err + } + } s.m[req.SenderUserID] = append(s.m[req.SenderUserID], sender) if req.SenderUserID != req.RecipientUserID && !req.RecipientBlocked { s.m[req.RecipientUserID] = append(s.m[req.RecipientUserID], recipient) } + if req.RandomID != 0 { + s.privateSendDedup[privateSendDedupKey{senderUserID: req.SenderUserID, randomID: req.RandomID}] = privateSendDedupRecord{ + recipientUserID: req.RecipientUserID, + senderSnapshot: senderSnapshot, + recipientMessage: immutablePrivateSendReceipt(recipient), + fingerprint: append([]byte(nil), fingerprint...), + } + } if s.dialogs != nil { if recipient.ID != 0 { s.upsertMemoryDialogsLocked(sender, recipient) @@ -128,6 +148,81 @@ func (s *MessageStore) SendPrivateText(_ context.Context, req domain.SendPrivate }, nil } +// LookupPrivateSendReplay returns an existing immutable/current replay receipt without running +// any send permission, reply resolution or allocation path. +func (s *MessageStore) LookupPrivateSendReplay(_ context.Context, req domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) { + if req.SenderUserID == 0 || req.RecipientUserID == 0 || req.RandomID == 0 { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("memory private send replay: invalid scope") + } + if err := store.ValidateSendFingerprint(req.IdempotencyFingerprint, "private send replay"); err != nil { + return domain.SendPrivateTextResult{}, false, err + } + s.mu.RLock() + defer s.mu.RUnlock() + return s.lookupPrivateSendReplayLocked(req) +} + +func (s *MessageStore) lookupPrivateSendReplayLocked(req domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) { + record, ok := s.privateSendDedup[privateSendDedupKey{senderUserID: req.SenderUserID, randomID: req.RandomID}] + if !ok { + return domain.SendPrivateTextResult{}, false, nil + } + if record.recipientUserID != req.RecipientUserID || !store.SameSendFingerprint(record.fingerprint, req.IdempotencyFingerprint) { + return domain.SendPrivateTextResult{}, false, domain.ErrMessageRandomIDDuplicate + } + firstSender, err := store.DecodePrivateSendSnapshot(record.senderSnapshot) + if err != nil { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("memory duplicate private message snapshot: %w", err) + } + sender := firstSender + visible := false + for _, current := range s.m[req.SenderUserID] { + if current.UID == firstSender.UID && current.ID == firstSender.ID { + sender = cloneMessage(current) + sender.RandomID = firstSender.RandomID + visible = true + break + } + } + if !visible && record.senderDeleteEvent == nil { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("memory duplicate private message %d is absent without a durable sender delete receipt", firstSender.UID) + } + recipient := cloneMessage(record.recipientMessage) + var replayDelete *domain.UpdateEvent + if record.senderDeleteEvent != nil { + cloned := cloneUpdateEvent(*record.senderDeleteEvent) + replayDelete = &cloned + } + return domain.SendPrivateTextResult{ + SenderMessage: sender, + RecipientMessage: recipient, + SenderEvent: newMessageEvent(firstSender), + RecipientEvent: newMessageEvent(recipient), + Duplicate: true, + ReplayDeleteEvent: replayDelete, + }, true, nil +} + +// immutablePrivateSendReceipt keeps the recipient allocation facts used by +// store-level idempotency tests. The sender response snapshot is stored as a +// versioned JSON value above so all nested media/reply graphs are immutable. +func immutablePrivateSendReceipt(msg domain.Message) domain.Message { + if msg.ID == 0 { + return domain.Message{} + } + return domain.Message{ + ID: msg.ID, + UID: msg.UID, + RandomID: msg.RandomID, + OwnerUserID: msg.OwnerUserID, + Peer: msg.Peer, + From: msg.From, + Date: msg.Date, + Out: msg.Out, + Pts: msg.Pts, + } +} + func (s *MessageStore) resolveMemoryReplyLocked(req domain.SendPrivateTextRequest) (*domain.MessageReply, *domain.MessageReply, error) { if req.ReplyTo == nil { return nil, nil, nil diff --git a/internal/store/memory/message_store.go b/internal/store/memory/message_store.go index aab44e49..152bc2a7 100644 --- a/internal/store/memory/message_store.go +++ b/internal/store/memory/message_store.go @@ -7,14 +7,17 @@ import ( // MessageStore 是 store.MessageStore 的内存实现。 type MessageStore struct { - mu sync.RWMutex - m map[int64][]domain.Message - nextUID int64 - nextBox map[int64]int - nextPts map[int64]int - readOutboxDates map[readOutboxDateKey]int - privateReactions map[int64]map[int64][]domain.ChannelMessagePeerReaction - dialogs *DialogStore + mu sync.RWMutex + m map[int64][]domain.Message + nextUID int64 + nextBox map[int64]int + nextPts map[int64]int + readOutboxDates map[readOutboxDateKey]int + privateReactions map[int64]map[int64][]domain.ChannelMessagePeerReaction + privateSendDedup map[privateSendDedupKey]privateSendDedupRecord + loginCodeDeliveries map[[32]byte]loginCodeDeliveryRecord + albumGroups map[albumGroupKey]albumGroupRecord + dialogs *DialogStore // polls 是共享 poll 权威（投票校验与读路径 enrichment）；nil 时 poll 链路按未接入处理。 polls *PollStore // savedPins 是收藏夹子会话置顶顺序（下标即 pinned_order，越小越前）。 @@ -35,13 +38,16 @@ type readOutboxDateKey struct { // NewMessageStore 创建内存 MessageStore。 func NewMessageStore(dialogs ...*DialogStore) *MessageStore { s := &MessageStore{ - m: make(map[int64][]domain.Message), - nextUID: 1, - nextBox: make(map[int64]int), - nextPts: make(map[int64]int), - readOutboxDates: make(map[readOutboxDateKey]int), - privateReactions: make(map[int64]map[int64][]domain.ChannelMessagePeerReaction), - savedPins: make(map[int64][]domain.Peer), + m: make(map[int64][]domain.Message), + nextUID: 1, + nextBox: make(map[int64]int), + nextPts: make(map[int64]int), + readOutboxDates: make(map[readOutboxDateKey]int), + privateReactions: make(map[int64]map[int64][]domain.ChannelMessagePeerReaction), + privateSendDedup: make(map[privateSendDedupKey]privateSendDedupRecord), + loginCodeDeliveries: make(map[[32]byte]loginCodeDeliveryRecord), + albumGroups: make(map[albumGroupKey]albumGroupRecord), + savedPins: make(map[int64][]domain.Peer), } if len(dialogs) > 0 { s.dialogs = dialogs[0] diff --git a/internal/store/memory/message_test.go b/internal/store/memory/message_test.go index e913c82a..6f265a25 100644 --- a/internal/store/memory/message_test.go +++ b/internal/store/memory/message_test.go @@ -162,7 +162,7 @@ func TestMessageStoreWebViewDataServiceActionRoundTrip(t *testing.T) { assertWebViewData("sender", got.SenderMessage) assertWebViewData("recipient", got.RecipientMessage) - dup, err := messages.SendPrivateText(ctx, domain.SendPrivateTextRequest{ + _, err = messages.SendPrivateText(ctx, domain.SendPrivateTextRequest{ SenderUserID: req.SenderUserID, RecipientUserID: req.RecipientUserID, RandomID: req.RandomID, @@ -178,13 +178,9 @@ func TestMessageStoreWebViewDataServiceActionRoundTrip(t *testing.T) { }, }, }) - if err != nil { - t.Fatalf("SendPrivateText duplicate: %v", err) + if !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("changed-media duplicate err = %v, want ErrMessageRandomIDDuplicate", err) } - if !dup.Duplicate || dup.SenderMessage.ID != got.SenderMessage.ID || dup.RecipientMessage.ID != got.RecipientMessage.ID { - t.Fatalf("duplicate = %+v, want original boxes", dup) - } - assertWebViewData("duplicate sender", dup.SenderMessage) recipientHistory, err := messages.ListByUser(ctx, req.RecipientUserID, domain.MessageFilter{ HasPeer: true, diff --git a/internal/store/memory/scoped_code_state_test.go b/internal/store/memory/scoped_code_state_test.go new file mode 100644 index 00000000..e3f05642 --- /dev/null +++ b/internal/store/memory/scoped_code_state_test.go @@ -0,0 +1,240 @@ +package memory + +import ( + "context" + "sync" + "testing" + "time" + + "telesrv/internal/store" +) + +func TestCodeStoreAtomicScopedVerification(t *testing.T) { + ctx := context.Background() + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016021", + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + Purpose: store.PhoneCodePurposeChangePhone, + UserID: 420021, + AuthKeyID: [8]byte{1, 2, 3, 4}, + MaxAttempts: 2, + } + } + + t.Run("only the active hash and exact scope can mutate", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord() + if err := codes.Set(ctx, "scoped-old", record, time.Minute); err != nil { + t.Fatal(err) + } + if err := codes.Set(ctx, "scoped-current", record, time.Minute); err != nil { + t.Fatal(err) + } + if result, err := codes.VerifyScoped(ctx, "scoped-old", record.Scope(), record.Code, 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("old-hash verify=%+v err=%v", result, err) + } + + otherScope := record.Scope() + otherScope.AuthKeyID = [8]byte{9} + if result, err := codes.VerifyScoped(ctx, "scoped-current", otherScope, "00000", 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("cross-scope verify=%+v err=%v", result, err) + } + stored, found, err := codes.Get(ctx, "scoped-current") + if err != nil || !found || stored.Attempts != 0 { + t.Fatalf("victim after cross-scope verify=%+v found=%v err=%v", stored, found, err) + } + }) + + t.Run("wrong attempts preserve ttl then delete code and index", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord() + if err := codes.Set(ctx, "scoped-wrong", record, time.Minute); err != nil { + t.Fatal(err) + } + before := codes.m["scoped-wrong"] + first, err := codes.VerifyScoped(ctx, "scoped-wrong", record.Scope(), "00000", 9) + if err != nil || first.Status != store.LoginCodeVerifyInvalid || first.Record.Attempts != 1 { + t.Fatalf("first wrong=%+v err=%v", first, err) + } + after := codes.m["scoped-wrong"] + if !after.expires.Equal(before.expires) || after.code.Revision == before.code.Revision { + t.Fatalf("wrong attempt expiry/revision before=%+v after=%+v", before, after) + } + second, err := codes.VerifyScoped(ctx, "scoped-wrong", record.Scope(), "00000", 9) + if err != nil || second.Status != store.LoginCodeVerifyInvalid || second.Record.Attempts != 2 { + t.Fatalf("threshold wrong=%+v err=%v", second, err) + } + if _, found, _ := codes.Get(ctx, "scoped-wrong"); found { + t.Fatal("threshold-exhausted scoped code remains") + } + if got := codes.scopes[record.Scope()]; got != "" { + t.Fatalf("threshold-exhausted scope index=%q, want missing", got) + } + if result, err := codes.VerifyScoped(ctx, "scoped-wrong", record.Scope(), record.Code, 9); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("verify after exhaustion=%+v err=%v", result, err) + } + }) + + t.Run("correct code consumes both keys exactly once", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord() + if err := codes.Set(ctx, "scoped-correct", record, time.Minute); err != nil { + t.Fatal(err) + } + expected, found, err := codes.Get(ctx, "scoped-correct") + if err != nil || !found { + t.Fatalf("get expected found=%v err=%v", found, err) + } + accepted, err := codes.VerifyScoped(ctx, "scoped-correct", record.Scope(), record.Code, 5) + if err != nil || accepted.Status != store.LoginCodeVerifyAccepted || accepted.Record != expected { + t.Fatalf("accepted=%+v err=%v, want %+v", accepted, err, expected) + } + if _, found, _ := codes.Get(ctx, "scoped-correct"); found || codes.scopes[record.Scope()] != "" { + t.Fatal("accepted scoped code or index remains") + } + if repeated, err := codes.VerifyScoped(ctx, "scoped-correct", record.Scope(), record.Code, 5); err != nil || repeated.Status != store.LoginCodeVerifyMissing { + t.Fatalf("repeated verify=%+v err=%v", repeated, err) + } + }) + + t.Run("legacy and inconsistent records fail closed", func(t *testing.T) { + codes := NewCodeStore() + legacy := newRecord() + legacy.Version = 0 + if err := codes.Set(ctx, "scoped-legacy", legacy, time.Minute); err != nil { + t.Fatal(err) + } + if result, err := codes.VerifyScoped(ctx, "scoped-legacy", legacy.Scope(), legacy.Code, 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("legacy verify=%+v err=%v", result, err) + } + if _, found, _ := codes.Get(ctx, "scoped-legacy"); found || codes.scopes[legacy.Scope()] != "" { + t.Fatal("legacy code or index remains") + } + + inconsistent := newRecord() + if err := codes.Set(ctx, "scoped-inconsistent", inconsistent, time.Minute); err != nil { + t.Fatal(err) + } + entry := codes.m["scoped-inconsistent"] + entry.code.Phone = "15550016999" + codes.m["scoped-inconsistent"] = entry + if result, err := codes.VerifyScoped(ctx, "scoped-inconsistent", inconsistent.Scope(), inconsistent.Code, 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("inconsistent verify=%+v err=%v", result, err) + } + if _, found, _ := codes.Get(ctx, "scoped-inconsistent"); found || codes.scopes[inconsistent.Scope()] != "" { + t.Fatal("inconsistent code or selecting index remains") + } + }) +} + +func TestCodeStoreAtomicScopedConcurrency(t *testing.T) { + ctx := context.Background() + const workers = 64 + newRecord := func(maxAttempts int) store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016022", + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + Purpose: store.PhoneCodePurposeChangePhone, + UserID: 420022, + AuthKeyID: [8]byte{5, 6, 7, 8}, + MaxAttempts: maxAttempts, + } + } + + t.Run("correct verification has one winner", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord(7) + if err := codes.Set(ctx, "scoped-verify-race", record, time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentMemoryScopedVerify(t, codes, "scoped-verify-race", record.Scope(), record.Code, workers) + if statuses[store.LoginCodeVerifyAccepted] != 1 || statuses[store.LoginCodeVerifyMissing] != workers-1 || statuses[store.LoginCodeVerifyInvalid] != 0 { + t.Fatalf("correct race statuses=%+v", statuses) + } + }) + + t.Run("wrong attempts cannot be lost", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord(7) + if err := codes.Set(ctx, "scoped-wrong-race", record, time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentMemoryScopedVerify(t, codes, "scoped-wrong-race", record.Scope(), "00000", workers) + if statuses[store.LoginCodeVerifyInvalid] != 7 || statuses[store.LoginCodeVerifyMissing] != workers-7 { + t.Fatalf("wrong race statuses=%+v", statuses) + } + if _, found, _ := codes.Get(ctx, "scoped-wrong-race"); found || codes.scopes[record.Scope()] != "" { + t.Fatal("wrong race left code or scope index") + } + }) + + t.Run("verification and cancellation share one winner", func(t *testing.T) { + codes := NewCodeStore() + record := newRecord(7) + if err := codes.Set(ctx, "scoped-mixed-race", record, time.Minute); err != nil { + t.Fatal(err) + } + results := make(chan bool, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(cancel bool) { + defer wg.Done() + if cancel { + _, found, err := codes.ConsumeScoped(ctx, "scoped-mixed-race", record.Scope()) + if err != nil { + t.Errorf("ConsumeScoped: %v", err) + } + results <- found + return + } + result, err := codes.VerifyScoped(ctx, "scoped-mixed-race", record.Scope(), record.Code, 5) + if err != nil { + t.Errorf("VerifyScoped: %v", err) + } + results <- result.Status == store.LoginCodeVerifyAccepted + }(i%2 == 0) + } + wg.Wait() + close(results) + winners := 0 + for won := range results { + if won { + winners++ + } + } + if winners != 1 { + t.Fatalf("verify/cancel winners=%d, want 1", winners) + } + }) +} + +func concurrentMemoryScopedVerify(t *testing.T, codes *CodeStore, hash string, scope store.PhoneCodeScope, code string, workers int) map[store.LoginCodeVerifyStatus]int { + t.Helper() + results := make(chan store.LoginCodeVerifyStatus, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + result, err := codes.VerifyScoped(context.Background(), hash, scope, code, 5) + if err != nil { + t.Errorf("VerifyScoped: %v", err) + return + } + results <- result.Status + }() + } + wg.Wait() + close(results) + statuses := make(map[store.LoginCodeVerifyStatus]int) + for status := range results { + statuses[status]++ + } + return statuses +} diff --git a/internal/store/memory/updates.go b/internal/store/memory/updates.go index 9465e34d..68d253c9 100644 --- a/internal/store/memory/updates.go +++ b/internal/store/memory/updates.go @@ -8,8 +8,9 @@ import ( // UpdateStateStore 是 store.UpdateStateStore 的内存实现。 type UpdateStateStore struct { - mu sync.RWMutex - states map[updateStateKey]domain.UpdateState + mu sync.RWMutex + states map[updateStateKey]domain.UpdateState + observed map[updateStateKey]domain.UpdateState } // UpdateEventStore 是 store.UpdateEventStore 的内存实现。 @@ -157,7 +158,10 @@ func (s *UpdateEventStore) MaxContiguousPts(_ context.Context, userID int64) (in // NewUpdateStateStore 创建内存 UpdateStateStore。 func NewUpdateStateStore() *UpdateStateStore { - return &UpdateStateStore{states: make(map[updateStateKey]domain.UpdateState)} + return &UpdateStateStore{ + states: make(map[updateStateKey]domain.UpdateState), + observed: make(map[updateStateKey]domain.UpdateState), + } } func (s *UpdateStateStore) Get(_ context.Context, id [8]byte, userID int64) (domain.UpdateState, bool, error) { @@ -191,9 +195,39 @@ func (s *UpdateStateStore) Save(_ context.Context, id [8]byte, userID int64, st return nil } +func (s *UpdateStateStore) ObserveClientState(_ context.Context, id [8]byte, userID int64, st domain.UpdateState) error { + s.mu.Lock() + key := updateStateKey{authKeyID: id, userID: userID} + prev := s.observed[key] + if st.Pts < prev.Pts { + st.Pts = prev.Pts + } + if st.Qts < prev.Qts { + st.Qts = prev.Qts + } + if st.Date < prev.Date { + st.Date = prev.Date + } + if st.Seq < prev.Seq { + st.Seq = prev.Seq + } + s.observed[key] = st + s.mu.Unlock() + return nil +} + +// ObservedClientState 暴露给同包/服务测试验证 retention 安全水位；业务读路径仍用 Get。 +func (s *UpdateStateStore) ObservedClientState(id [8]byte, userID int64) (domain.UpdateState, bool) { + s.mu.RLock() + defer s.mu.RUnlock() + st, ok := s.observed[updateStateKey{authKeyID: id, userID: userID}] + return st, ok +} + func (s *UpdateStateStore) Delete(_ context.Context, id [8]byte, userID int64) error { s.mu.Lock() delete(s.states, updateStateKey{authKeyID: id, userID: userID}) + delete(s.observed, updateStateKey{authKeyID: id, userID: userID}) s.mu.Unlock() return nil } @@ -203,6 +237,7 @@ func (s *UpdateStateStore) DeleteAuthKey(_ context.Context, id [8]byte) error { for k := range s.states { if k.authKeyID == id { delete(s.states, k) + delete(s.observed, k) } } s.mu.Unlock() diff --git a/internal/store/postgres/album_group.go b/internal/store/postgres/album_group.go new file mode 100644 index 00000000..6e8d4f38 --- /dev/null +++ b/internal/store/postgres/album_group.go @@ -0,0 +1,162 @@ +package postgres + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/binary" + "errors" + "fmt" + "sort" + + "telesrv/internal/domain" +) + +// ReserveAlbumGroup 先按稳定顺序获取整批 key 的事务级 advisory locks，再读取旧绑定 +// 并一次性补齐缺失项。锁覆盖不存在的行，因此避免单靠 UNIQUE/ON CONFLICT 时两个实例 +// 对重叠批次分别选出不同 grouped_id 的 write-skew。 +func (s *MessageStore) ReserveAlbumGroup(ctx context.Context, req domain.AlbumGroupReservationRequest) (int64, error) { + if err := req.Validate(); err != nil { + return 0, err + } + beginner, ok := s.db.(txBeginner) + if !ok { + return 0, errors.New("reserve album group requires transaction-capable postgres handle") + } + tx, err := beginner.Begin(ctx) + if err != nil { + return 0, fmt.Errorf("reserve album group begin: %w", err) + } + defer func() { _ = tx.Rollback(ctx) }() + + lockIDs := albumGroupAdvisoryLockIDs(req) + for _, lockID := range lockIDs { + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1)`, lockID); err != nil { + return 0, fmt.Errorf("reserve album group lock: %w", err) + } + } + + randomIDs := make([]int64, 0, len(req.Items)) + requestedIntents := make(map[int64][]byte, len(req.Items)) + for _, item := range req.Items { + randomIDs = append(randomIDs, item.RandomID) + requestedIntents[item.RandomID] = item.IntentHash + } + rows, err := tx.Query(ctx, ` +SELECT random_id, grouped_id, intent_hash +FROM album_group_reservations +WHERE sender_user_id = $1 + AND peer_type = $2 + AND peer_id = $3 + AND random_id = ANY($4::bigint[]) +ORDER BY random_id`, req.SenderUserID, string(req.Peer.Type), req.Peer.ID, randomIDs) + if err != nil { + return 0, fmt.Errorf("reserve album group read existing: %w", err) + } + existingGroups := make(map[int64]struct{}, 2) + for rows.Next() { + var randomID int64 + var groupedID int64 + var intentHash []byte + if err := rows.Scan(&randomID, &groupedID, &intentHash); err != nil { + rows.Close() + return 0, fmt.Errorf("reserve album group scan existing: %w", err) + } + if !bytes.Equal(intentHash, requestedIntents[randomID]) { + rows.Close() + return 0, fmt.Errorf("%w: album random_id %d intent changed", domain.ErrMessageRandomIDDuplicate, randomID) + } + existingGroups[groupedID] = struct{}{} + } + readErr := rows.Err() + rows.Close() + if readErr != nil { + return 0, fmt.Errorf("reserve album group iterate existing: %w", readErr) + } + if len(existingGroups) > 1 { + return 0, fmt.Errorf("%w: album request spans multiple grouped_id values", domain.ErrMessageRandomIDDuplicate) + } + + groupedID := req.ProposedGroupedID + for existingGroup := range existingGroups { + groupedID = existingGroup + } + for _, item := range req.Items { + if _, err := tx.Exec(ctx, ` +INSERT INTO album_group_reservations ( + sender_user_id, peer_type, peer_id, random_id, intent_hash, grouped_id +) +VALUES ($1, $2, $3, $4, $5, $6) +ON CONFLICT (sender_user_id, peer_type, peer_id, random_id) DO NOTHING`, + req.SenderUserID, string(req.Peer.Type), req.Peer.ID, item.RandomID, item.IntentHash, groupedID); err != nil { + return 0, fmt.Errorf("reserve album group insert binding: %w", err) + } + } + + // 防御性复核：advisory lock 协议若被未来代码绕过，也不能把拆组状态作为成功返回。 + rows, err = tx.Query(ctx, ` +SELECT random_id, grouped_id, intent_hash +FROM album_group_reservations +WHERE sender_user_id = $1 + AND peer_type = $2 + AND peer_id = $3 + AND random_id = ANY($4::bigint[])`, + req.SenderUserID, string(req.Peer.Type), req.Peer.ID, randomIDs) + if err != nil { + return 0, fmt.Errorf("reserve album group verify: %w", err) + } + verified := 0 + for rows.Next() { + var randomID, storedGroup int64 + var intentHash []byte + if err := rows.Scan(&randomID, &storedGroup, &intentHash); err != nil { + rows.Close() + return 0, fmt.Errorf("reserve album group verify scan: %w", err) + } + if storedGroup != groupedID || !bytes.Equal(intentHash, requestedIntents[randomID]) { + rows.Close() + return 0, fmt.Errorf("%w: album reservation diverged for random_id %d", domain.ErrMessageRandomIDDuplicate, randomID) + } + verified++ + } + verifyErr := rows.Err() + rows.Close() + if verifyErr != nil { + return 0, fmt.Errorf("reserve album group verify iterate: %w", verifyErr) + } + if verified != len(req.Items) { + return 0, fmt.Errorf("%w: album reservation count=%d/%d", domain.ErrMessageRandomIDDuplicate, verified, len(req.Items)) + } + if err := tx.Commit(ctx); err != nil { + return 0, fmt.Errorf("reserve album group commit: %w", err) + } + return groupedID, nil +} + +// albumGroupAdvisoryLockIDs 对每个业务 key 派生一个 64-bit advisory lock，并按数值 +// 排序、去重。hash 碰撞最多造成无害串行化；排序保证重叠批次不会互相反序死锁。 +func albumGroupAdvisoryLockIDs(req domain.AlbumGroupReservationRequest) []int64 { + ids := make([]int64, 0, len(req.Items)) + for _, item := range req.Items { + h := sha256.New() + _, _ = h.Write([]byte("telesrv:album-group:v1\x00")) + var word [8]byte + binary.BigEndian.PutUint64(word[:], uint64(req.SenderUserID)) + _, _ = h.Write(word[:]) + _, _ = h.Write([]byte(req.Peer.Type)) + binary.BigEndian.PutUint64(word[:], uint64(req.Peer.ID)) + _, _ = h.Write(word[:]) + binary.BigEndian.PutUint64(word[:], uint64(item.RandomID)) + _, _ = h.Write(word[:]) + sum := h.Sum(nil) + ids = append(ids, int64(binary.BigEndian.Uint64(sum[:8]))) + } + sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) + out := ids[:0] + for _, id := range ids { + if len(out) == 0 || out[len(out)-1] != id { + out = append(out, id) + } + } + return out +} diff --git a/internal/store/postgres/album_group_integration_test.go b/internal/store/postgres/album_group_integration_test.go new file mode 100644 index 00000000..efd956ce --- /dev/null +++ b/internal/store/postgres/album_group_integration_test.go @@ -0,0 +1,118 @@ +package postgres + +import ( + "context" + "crypto/sha256" + "errors" + "sync" + "testing" + + "telesrv/internal/domain" +) + +func pgAlbumItem(randomID int64, label string) domain.AlbumGroupReservationItem { + sum := sha256.Sum256([]byte(label)) + return domain.AlbumGroupReservationItem{RandomID: randomID, IntentHash: sum[:]} +} + +func pgAlbumReq(sender int64, peer domain.Peer, group int64, items ...domain.AlbumGroupReservationItem) domain.AlbumGroupReservationRequest { + return domain.AlbumGroupReservationRequest{ + SenderUserID: sender, + Peer: peer, + Items: items, + ProposedGroupedID: group, + } +} + +func TestAlbumGroupReservationConvergesAcrossPostgresInstances(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + sender, err := users.Create(ctx, domain.User{AccessHash: 7601, Phone: "+1760" + suffix + "01", FirstName: "AlbumSender"}) + if err != nil { + t.Fatalf("create sender: %v", err) + } + recipient, err := users.Create(ctx, domain.User{AccessHash: 7602, Phone: "+1760" + suffix + "02", FirstName: "AlbumRecipient"}) + if err != nil { + t.Fatalf("create recipient: %v", err) + } + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM album_group_reservations WHERE sender_user_id = $1", sender.ID) + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{sender.ID, recipient.ID}) + }) + + privatePeer := domain.Peer{Type: domain.PeerTypeUser, ID: recipient.ID} + full := []domain.AlbumGroupReservationItem{ + pgAlbumItem(76001, "one"), + pgAlbumItem(76002, "two"), + pgAlbumItem(76003, "three"), + } + firstStore := NewMessageStore(pool) + groupedID, err := firstStore.ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 761, full...)) + if err != nil || groupedID != 761 { + t.Fatalf("reserve full = %d err=%v, want 761", groupedID, err) + } + // 新 store 实例模拟另一进程；失败子集必须恢复首次整包的组。 + replayed, err := NewMessageStore(pool).ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 762, full[1:]...)) + if err != nil || replayed != groupedID { + t.Fatalf("reserve subset = %d err=%v, want %d", replayed, err, groupedID) + } + if _, err := NewMessageStore(pool).ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 763, pgAlbumItem(76002, "changed"))); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("changed intent err=%v, want ErrMessageRandomIDDuplicate", err) + } + + // 同 random_id 在不同 peer 作用域独立，不会误并相册。 + channelPeer := domain.Peer{Type: domain.PeerTypeChannel, ID: recipient.ID} + channelGroup, err := firstStore.ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, channelPeer, 764, full[0])) + if err != nil || channelGroup != 764 { + t.Fatalf("channel peer isolated group = %d err=%v, want 764", channelGroup, err) + } + + // 两个实例同时预留部分重叠的批次，shared random_id 的 advisory lock 必须让 + // 两边串行收敛；最终 4/5/6 三个 item 全部同组。 + left := pgAlbumReq(sender.ID, privatePeer, 765, pgAlbumItem(76004, "four"), pgAlbumItem(76005, "shared")) + right := pgAlbumReq(sender.ID, privatePeer, 766, pgAlbumItem(76005, "shared"), pgAlbumItem(76006, "six")) + requests := []domain.AlbumGroupReservationRequest{left, right} + results := make([]int64, 2) + errs := make([]error, 2) + start := make(chan struct{}) + var wg sync.WaitGroup + for i := range requests { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + results[i], errs[i] = NewMessageStore(pool).ReserveAlbumGroup(ctx, requests[i]) + }(i) + } + close(start) + wg.Wait() + if errs[0] != nil || errs[1] != nil || results[0] == 0 || results[0] != results[1] { + t.Fatalf("concurrent groups=%v errs=%v, want one non-zero group", results, errs) + } + for _, item := range []domain.AlbumGroupReservationItem{left.Items[0], left.Items[1], right.Items[1]} { + got, err := NewMessageStore(pool).ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 767, item)) + if err != nil || got != results[0] { + t.Fatalf("verify random_id %d = %d err=%v, want %d", item.RandomID, got, err, results[0]) + } + } + + // 一个请求同时命中两个历史组必须整体失败，不能绑定其中的新 item。 + oldA := pgAlbumItem(76007, "old-a") + oldB := pgAlbumItem(76008, "old-b") + newItem := pgAlbumItem(76009, "new") + if _, err := firstStore.ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 768, oldA)); err != nil { + t.Fatal(err) + } + if _, err := firstStore.ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 769, oldB)); err != nil { + t.Fatal(err) + } + if _, err := firstStore.ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 770, oldA, oldB, newItem)); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("mixed old groups err=%v, want ErrMessageRandomIDDuplicate", err) + } + fresh, err := firstStore.ReserveAlbumGroup(ctx, pgAlbumReq(sender.ID, privatePeer, 771, newItem)) + if err != nil || fresh != 771 { + t.Fatalf("post-conflict fresh item = %d err=%v, want 771", fresh, err) + } +} diff --git a/internal/store/postgres/auth_login_integration_test.go b/internal/store/postgres/auth_login_integration_test.go index c075da8c..4754b094 100644 --- a/internal/store/postgres/auth_login_integration_test.go +++ b/internal/store/postgres/auth_login_integration_test.go @@ -34,6 +34,7 @@ func TestAuthSignUpWritesOfficialLoginMessagePostgres(t *testing.T) { nil, "12345", appauth.WithLoginMessages(messages, dialogs), + appauth.WithLoginCodeDelivery(messages), ) var authKeyID [8]byte @@ -85,7 +86,7 @@ func TestAuthSignUpWritesOfficialLoginMessagePostgres(t *testing.T) { } } -func TestAuthSignInOfficialLoginMessagePreservesReadWatermarkPostgres(t *testing.T) { +func TestAuthSendCodeOfficialLoginMessagePreservesReadWatermarkBeforeSignInPostgres(t *testing.T) { pool := testPool(t) ctx := context.Background() @@ -105,6 +106,7 @@ func TestAuthSignInOfficialLoginMessagePreservesReadWatermarkPostgres(t *testing nil, "12345", appauth.WithLoginMessages(messages, dialogs), + appauth.WithLoginCodeDelivery(messages), ) var authKeyID [8]byte @@ -126,6 +128,9 @@ func TestAuthSignInOfficialLoginMessagePreservesReadWatermarkPostgres(t *testing if err != nil { t.Fatalf("SendCode signup: %v", err) } + if _, _, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: authKeyID}, phone, hash, "12345"); err != nil || !needSignUp { + t.Fatalf("SignIn before signup needSignUp=%v err=%v, want true/nil", needSignUp, err) + } u, first, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: authKeyID}, phone, hash, "PgLogin", "Read") if err != nil { t.Fatalf("SignUp: %v", err) @@ -180,19 +185,29 @@ FROM target`, u.ID, domain.OfficialSystemUserID).Scan(&top, &readMax, &unread, & if err != nil { t.Fatalf("SendCode signin second: %v", err) } - _, second, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: authKeyID}, phone, hash, "12345") + secondID := first.ID + 1 + assertOfficialDialog(secondID, first.ID, 1) + _, lateSecond, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: authKeyID}, phone, hash, "12345") if err != nil || needSignUp { t.Fatalf("SignIn second needSignUp=%v err=%v", needSignUp, err) } - assertOfficialDialog(second.ID, first.ID, 1) + if lateSecond.ID != 0 { + t.Fatalf("SignIn second returned late login message %+v, want zero", lateSecond) + } + assertOfficialDialog(secondID, 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{AuthKeyID: authKeyID}, phone, hash, "12345") + thirdID := first.ID + 2 + assertOfficialDialog(thirdID, first.ID, 2) + _, lateThird, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: authKeyID}, phone, hash, "12345") if err != nil || needSignUp { t.Fatalf("SignIn third needSignUp=%v err=%v", needSignUp, err) } - assertOfficialDialog(third.ID, first.ID, 2) + if lateThird.ID != 0 { + t.Fatalf("SignIn third returned late login message %+v, want zero", lateThird) + } + assertOfficialDialog(thirdID, first.ID, 2) } diff --git a/internal/store/postgres/authkey.go b/internal/store/postgres/authkey.go index b78ee951..e0af8c25 100644 --- a/internal/store/postgres/authkey.go +++ b/internal/store/postgres/authkey.go @@ -5,6 +5,7 @@ import ( "encoding/binary" "errors" "fmt" + "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgtype" @@ -25,20 +26,23 @@ func NewAuthKeyStore(db sqlcgen.DBTX) *AuthKeyStore { } // Save 实现 store.AuthKeyStore。auth_key_id 以小端解释为 int64 存入 BIGINT； -// created_at 交由 DB 默认值（now()），故传入的 CreatedAt 不落库。 +// created_at/last_used_at 交由 DB 默认值（now()），故传入的 CreatedAt 不落库。 func (s *AuthKeyStore) Save(ctx context.Context, k store.AuthKeyData) error { if _, err := s.db.Exec(ctx, ` INSERT INTO auth_keys (auth_key_id, body, server_salt) VALUES ($1, $2, $3) ON CONFLICT (auth_key_id) DO UPDATE -SET body = EXCLUDED.body, server_salt = EXCLUDED.server_salt +SET body = EXCLUDED.body, server_salt = EXCLUDED.server_salt, last_used_at = now() `, authKeyIDToInt64(k.ID), k.Value[:], k.ServerSalt); err != nil { return fmt.Errorf("upsert auth key: %w", err) } return nil } -// Get 实现 store.AuthKeyStore。不存在时 found=false。 +// Get 实现 store.AuthKeyStore。不存在时 found=false。读取与 last_used_at touch 是同一条 +// UPDATE ... RETURNING：若 orphan GC 已锁定并删除该行，Get 等待后得到 no rows；若 Get 先 +// 完成，GC 的 cutoff/final predicate 会看到新水位并跳过。这样连接不会在“读到旧 key、尚未 +// 注册进 SessionManager”的窗口被后台清理。 func (s *AuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) { var ( body []byte @@ -52,10 +56,11 @@ func (s *AuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, appVersion string ) err := s.db.QueryRow(ctx, ` -SELECT auth_key_id, body, server_salt, created_at, - layer, device_model, platform, system_version, api_id, app_version -FROM auth_keys +UPDATE auth_keys +SET last_used_at = now() WHERE auth_key_id = $1 +RETURNING auth_key_id, body, server_salt, created_at, + layer, device_model, platform, system_version, api_id, app_version `, authKeyIDToInt64(id)).Scan(new(int64), &body, &serverSalt, &createdAt, &layer, &deviceModel, &platform, &systemVersion, &apiID, &appVersion) if err != nil { if errors.Is(err, pgx.ErrNoRows) { @@ -83,6 +88,52 @@ WHERE auth_key_id = $1 return data, true, nil } +const activeAuthKeyHeartbeatBatch = 4096 + +// TouchActiveRawAuthKeys refreshes the durable activity lease for raw auth keys currently held by +// this server instance. Orphan collection is database-global while SessionManager is process-local; +// without this heartbeat, instance A can collect a long-lived unauthorised key that is active on +// instance B after its one-time Get touch ages past the retention cutoff. +// +// The caller runs this well inside the orphan-retention window and skips collection if a heartbeat +// fails. Batching keeps the ANY array and one UPDATE bounded at large connection counts. +func (s *AuthKeyStore) TouchActiveRawAuthKeys(ctx context.Context, ids [][8]byte) error { + if len(ids) == 0 { + return nil + } + seen := make(map[int64]struct{}, len(ids)) + keyIDs := make([]int64, 0, len(ids)) + for _, id := range ids { + keyID := authKeyIDToInt64(id) + if _, duplicate := seen[keyID]; duplicate { + continue + } + seen[keyID] = struct{}{} + keyIDs = append(keyIDs, keyID) + } + for start := 0; start < len(keyIDs); start += activeAuthKeyHeartbeatBatch { + end := start + activeAuthKeyHeartbeatBatch + if end > len(keyIDs) { + end = len(keyIDs) + } + batch := keyIDs[start:end] + tag, err := s.db.Exec(ctx, ` +UPDATE auth_keys +SET last_used_at = now() +WHERE auth_key_id = ANY($1::bigint[])`, batch) + if err != nil { + return fmt.Errorf("touch active raw auth keys: %w", err) + } + if tag.RowsAffected() != int64(len(batch)) { + return fmt.Errorf( + "touch active raw auth keys: refreshed %d of %d keys", + tag.RowsAffected(), len(batch), + ) + } + } + return nil +} + func (s *AuthKeyStore) UpdateClientInfo(ctx context.Context, id [8]byte, info store.AuthKeyClientInfo) error { if _, err := s.db.Exec(ctx, ` UPDATE auth_keys @@ -108,23 +159,100 @@ WHERE auth_key_id = $1 // raw temp key 重连时仍能进入 RPC 层，只得到 AUTH_KEY_UNREGISTERED，而不是连接层 404。 func (s *AuthKeyStore) Delete(ctx context.Context, id [8]byte) error { keyID := authKeyIDToInt64(id) - if _, err := s.db.Exec(ctx, ` -WITH doomed_temp AS ( + var touched int + if err := s.db.QueryRow(ctx, ` +WITH doomed_temp AS MATERIALIZED ( SELECT temp_auth_key_id FROM temp_auth_key_bindings WHERE perm_auth_key_id = $1 +), doomed_keys AS MATERIALIZED ( + SELECT $1::bigint AS auth_key_id + UNION + SELECT temp_auth_key_id FROM doomed_temp +), deleted_update_states AS ( + -- update_states intentionally has no auth_keys FK: remove device cursors in the + -- same statement/transaction as both the permanent and derived temp keys. + DELETE FROM update_states + WHERE auth_key_id IN (SELECT auth_key_id FROM doomed_keys) + RETURNING auth_key_id ), deleted_temp AS ( DELETE FROM auth_keys - WHERE auth_key_id IN (SELECT temp_auth_key_id FROM doomed_temp) + WHERE auth_key_id IN (SELECT auth_key_id FROM doomed_keys) + RETURNING auth_key_id ) -DELETE FROM auth_keys -WHERE auth_key_id = $1 -`, keyID); err != nil { +SELECT + (SELECT count(*) FROM deleted_update_states)::int + + (SELECT count(*) FROM deleted_temp)::int`, keyID).Scan(&touched); err != nil { return fmt.Errorf("delete auth key and temp bindings: %w", err) } return nil } +// DeleteOrphaned 回收握手已落库、但从未形成 authorization/temp binding 且当前没有 +// 活跃物理连接的旧 auth key。last_used_at 与 Get 的 UPDATE ... RETURNING 行锁配对，封住 +// active-key 快照之后新连接开始使用旧 key 的竞态；所有引用条件仍在最终 DELETE 中复核。 +// protected 必须是 SessionManager 的 raw key。 +func (s *AuthKeyStore) DeleteOrphaned(ctx context.Context, olderThan time.Duration, limit int, protected [][8]byte) (int, error) { + if olderThan <= 0 || limit <= 0 { + return 0, nil + } + if limit > 100000 { + limit = 100000 + } + protectedIDs := make([]int64, 0, len(protected)) + for _, id := range protected { + protectedIDs = append(protectedIDs, authKeyIDToInt64(id)) + } + var deleted int + err := s.db.QueryRow(ctx, ` +WITH candidates AS MATERIALIZED ( + SELECT k.auth_key_id + FROM auth_keys k + WHERE k.last_used_at < now() - make_interval(secs => $1::double precision) + AND NOT (k.auth_key_id = ANY($2::bigint[])) + AND NOT EXISTS ( + SELECT 1 FROM authorizations a WHERE a.auth_key_id = k.auth_key_id + ) + AND NOT EXISTS ( + SELECT 1 + FROM temp_auth_key_bindings b + WHERE b.temp_auth_key_id = k.auth_key_id OR b.perm_auth_key_id = k.auth_key_id + ) + ORDER BY k.last_used_at ASC, k.auth_key_id ASC + LIMIT $3 + FOR UPDATE OF k SKIP LOCKED +), deleted_update_states AS ( + -- Historical authorization-only deletion could leave a cursor without an + -- auth_keys FK. GC owns that stale row once the raw key is proven orphaned. + DELETE FROM update_states s + USING candidates c + WHERE s.auth_key_id = c.auth_key_id + RETURNING s.auth_key_id +), deleted_keys AS ( + DELETE FROM auth_keys k + USING candidates c + WHERE k.auth_key_id = c.auth_key_id + AND k.last_used_at < now() - make_interval(secs => $1::double precision) + AND NOT (k.auth_key_id = ANY($2::bigint[])) + AND NOT EXISTS ( + SELECT 1 FROM authorizations a WHERE a.auth_key_id = k.auth_key_id + ) + AND NOT EXISTS ( + SELECT 1 + FROM temp_auth_key_bindings b + WHERE b.temp_auth_key_id = k.auth_key_id OR b.perm_auth_key_id = k.auth_key_id + ) + RETURNING k.auth_key_id +) +SELECT count(*)::int +FROM deleted_keys +CROSS JOIN LATERAL (SELECT count(*) FROM deleted_update_states) AS touched`, olderThan.Seconds(), protectedIDs, limit).Scan(&deleted) + if err != nil { + return 0, fmt.Errorf("delete orphaned auth keys: %w", err) + } + return deleted, nil +} + // authKeyIDToInt64 把 [8]byte 的 auth_key_id 按小端解释为 int64（MTProto 定义即 SHA1 低 64 位）。 func authKeyIDToInt64(id [8]byte) int64 { return int64(binary.LittleEndian.Uint64(id[:])) diff --git a/internal/store/postgres/authkey_retention_integration_test.go b/internal/store/postgres/authkey_retention_integration_test.go new file mode 100644 index 00000000..ef3dcfaf --- /dev/null +++ b/internal/store/postgres/authkey_retention_integration_test.go @@ -0,0 +1,215 @@ +package postgres + +import ( + "context" + "crypto/rand" + "testing" + "time" + + "telesrv/internal/domain" + "telesrv/internal/store" +) + +func TestAuthKeyStoreDeleteOrphanedIsBoundedAndProtectsReferencesPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + userID := createRevokeTestUser(t, ctx, pool, "orphan-auth-key") + + newKey := func() [8]byte { + var id [8]byte + if _, err := rand.Read(id[:]); err != nil { + t.Fatalf("random auth key id: %v", err) + } + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key %x: %v", id, err) + } + t.Cleanup(func() { _ = keys.Delete(ctx, id) }) + return id + } + orphanOne, orphanTwo := newKey(), newKey() + recent := newKey() + authorized := newKey() + temp, perm := newKey(), newKey() + active := newKey() + if _, err := pool.Exec(ctx, ` +INSERT INTO update_states (auth_key_id, user_id, pts, observed_pts) +VALUES ($1, $3, 0, 0), ($2, $3, 0, 0)`, + authKeyIDToInt64(orphanOne), authKeyIDToInt64(orphanTwo), userID); err != nil { + t.Fatalf("insert stale orphan update states: %v", err) + } + + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: authorized, UserID: userID}); err != nil { + t.Fatalf("bind authorization: %v", err) + } + if err := NewTempAuthKeyBindingStore(pool).Save(ctx, domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, PermAuthKeyID: authKeyIDToInt64(perm), Nonce: 1, + TempSessionID: 2, ExpiresAt: int(time.Now().Add(time.Hour).Unix()), EncryptedMessage: []byte{1}, + }); err != nil { + t.Fatalf("save temp binding: %v", err) + } + + // Use a test-only historical window so a shared developer database's unrelated 24h-old + // handshake keys cannot win the bounded candidate slot or be mutated by this test. + const retention = 150 * 365 * 24 * time.Hour + old := time.Now().Add(-200 * 365 * 24 * time.Hour) + oldIDs := [][8]byte{orphanOne, orphanTwo, authorized, temp, perm, active} + for _, id := range oldIDs { + if _, err := pool.Exec(ctx, "UPDATE auth_keys SET created_at = $2, last_used_at = $2 WHERE auth_key_id = $1", authKeyIDToInt64(id), old); err != nil { + t.Fatalf("age auth key %x: %v", id, err) + } + } + + deleted, err := keys.DeleteOrphaned(ctx, retention, 1, [][8]byte{active}) + if err != nil || deleted != 1 { + t.Fatalf("first bounded orphan delete = %d/%v, want 1/nil", deleted, err) + } + var remainingOrphans int + if err := pool.QueryRow(ctx, ` +SELECT count(*) FROM auth_keys WHERE auth_key_id = ANY($1::bigint[]) +`, []int64{authKeyIDToInt64(orphanOne), authKeyIDToInt64(orphanTwo)}).Scan(&remainingOrphans); err != nil { + t.Fatalf("count remaining orphans: %v", err) + } + if remainingOrphans != 1 { + t.Fatalf("remaining old unreferenced orphans = %d, want 1 after batch=1", remainingOrphans) + } + + deleted, err = keys.DeleteOrphaned(ctx, retention, 20, [][8]byte{active}) + if err != nil || deleted != 1 { + t.Fatalf("second orphan delete = %d/%v, want remaining 1/nil", deleted, err) + } + var orphanStates int + if err := pool.QueryRow(ctx, ` +SELECT count(*)::int +FROM update_states +WHERE auth_key_id = ANY($1::bigint[])`, []int64{ + authKeyIDToInt64(orphanOne), authKeyIDToInt64(orphanTwo), + }).Scan(&orphanStates); err != nil { + t.Fatalf("count orphan update states: %v", err) + } + if orphanStates != 0 { + t.Fatalf("orphan update states = %d, want 0 after atomic key GC", orphanStates) + } + for name, id := range map[string][8]byte{ + "recent": recent, "authorized": authorized, "temp": temp, "perm": perm, "active": active, + } { + if _, found, err := keys.Get(ctx, id); err != nil || !found { + t.Fatalf("protected %s key %x found=%v err=%v, want retained", name, id, found, err) + } + } +} + +func TestAuthKeyStoreDeleteCleansPermanentAndTempUpdateStatesPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + userID := createRevokeTestUser(t, ctx, pool, "auth-key-delete-state") + perm := randomUpdateRetentionAuthKey(t) + temp := randomUpdateRetentionAuthKey(t) + for _, id := range [][8]byte{perm, temp} { + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key %x: %v", id, err) + } + id := id + t.Cleanup(func() { _ = keys.Delete(ctx, id) }) + } + if err := NewTempAuthKeyBindingStore(pool).Save(ctx, domain.TempAuthKeyBinding{ + TempAuthKeyID: temp, + PermAuthKeyID: authKeyIDToInt64(perm), + Nonce: 31, + TempSessionID: 32, + ExpiresAt: int(time.Now().Add(time.Hour).Unix()), + EncryptedMessage: []byte{1}, + }); err != nil { + t.Fatalf("save temp binding: %v", err) + } + if _, err := pool.Exec(ctx, ` +INSERT INTO update_states (auth_key_id, user_id, pts, observed_pts) +VALUES ($1, $3, 0, 0), ($2, $3, 0, 0)`, + authKeyIDToInt64(perm), authKeyIDToInt64(temp), userID); err != nil { + t.Fatalf("insert permanent/temp update states: %v", err) + } + + if err := keys.Delete(ctx, perm); err != nil { + t.Fatalf("delete permanent auth key: %v", err) + } + ids := []int64{authKeyIDToInt64(perm), authKeyIDToInt64(temp)} + var keyRows, stateRows int + if err := pool.QueryRow(ctx, ` +SELECT + (SELECT count(*) FROM auth_keys WHERE auth_key_id = ANY($1::bigint[]))::int, + (SELECT count(*) FROM update_states WHERE auth_key_id = ANY($1::bigint[]))::int`, ids).Scan(&keyRows, &stateRows); err != nil { + t.Fatalf("count deleted auth key state: %v", err) + } + if keyRows != 0 || stateRows != 0 { + t.Fatalf("remaining key/state rows = %d/%d, want 0/0", keyRows, stateRows) + } +} + +func TestAuthKeyGetTouchPreventsOrphanCollectionPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + var id [8]byte + if _, err := rand.Read(id[:]); err != nil { + t.Fatalf("random auth key id: %v", err) + } + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key: %v", err) + } + t.Cleanup(func() { _ = keys.Delete(ctx, id) }) + + const retention = 150 * 365 * 24 * time.Hour + old := time.Now().Add(-200 * 365 * 24 * time.Hour) + if _, err := pool.Exec(ctx, "UPDATE auth_keys SET created_at = $2, last_used_at = $2 WHERE auth_key_id = $1", authKeyIDToInt64(id), old); err != nil { + t.Fatalf("age auth key: %v", err) + } + if _, found, err := keys.Get(ctx, id); err != nil || !found { + t.Fatalf("touch auth key found=%v err=%v", found, err) + } + deleted, err := keys.DeleteOrphaned(ctx, retention, 10, nil) + if err != nil { + t.Fatalf("delete orphaned: %v", err) + } + if deleted != 0 { + t.Fatalf("deleted = %d, want 0 after atomic Get touch", deleted) + } + if _, found, err := keys.Get(ctx, id); err != nil || !found { + t.Fatalf("touched key retained found=%v err=%v", found, err) + } +} + +func TestActiveRawAuthKeyHeartbeatProtectsOtherInstanceKeyPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + keys := NewAuthKeyStore(pool) + var id [8]byte + if _, err := rand.Read(id[:]); err != nil { + t.Fatalf("random auth key id: %v", err) + } + t.Cleanup(func() { _ = keys.Delete(ctx, id) }) + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key: %v", err) + } + old := time.Now().Add(-48 * time.Hour) + if _, err := pool.Exec(ctx, "UPDATE auth_keys SET created_at = $2, last_used_at = $2 WHERE auth_key_id = $1", authKeyIDToInt64(id), old); err != nil { + t.Fatalf("age active auth key: %v", err) + } + + // Model another process heartbeating its local SessionManager snapshot. The collector on this + // process has no protected-list entry for the key and must still respect durable last_used_at. + if err := keys.TouchActiveRawAuthKeys(ctx, [][8]byte{id, id}); err != nil { + t.Fatalf("heartbeat active raw auth key: %v", err) + } + deleted, err := keys.DeleteOrphaned(ctx, 24*time.Hour, 10, nil) + if err != nil { + t.Fatalf("delete orphaned after heartbeat: %v", err) + } + if deleted != 0 { + t.Fatalf("deleted = %d, want active key protected by durable heartbeat", deleted) + } + if _, found, err := keys.Get(ctx, id); err != nil || !found { + t.Fatalf("heartbeat key found=%v err=%v, want present", found, err) + } +} diff --git a/internal/store/postgres/authorization.go b/internal/store/postgres/authorization.go index 819cd2cd..43481592 100644 --- a/internal/store/postgres/authorization.go +++ b/internal/store/postgres/authorization.go @@ -28,7 +28,111 @@ func (s *AuthorizationStore) Bind(ctx context.Context, a domain.Authorization) e if a.Hash == 0 { a.Hash = authorizationHash(a.AuthKeyID) } - _, err := s.db.Exec(ctx, ` + bind := func(db sqlcgen.DBTX) error { + return bindAuthorization(ctx, db, a) + } + var err error + if tx, ok := s.db.(pgx.Tx); ok { + err = bind(tx) + } else { + err = withTx(ctx, s.db, "bind authorization", func(tx pgx.Tx) error { + return bind(tx) + }) + } + if err != nil { + return fmt.Errorf("upsert authorization: %w", err) + } + return nil +} + +// bindAuthorization 把 auth_key→user 绑定和设备 update baseline 作为同一个状态边界提交。 +// +// 锁顺序固定为：auth_keys 母行 → 目标 user_update_watermarks → +// user_update_retention → 目标 update_states。前两个 user 锁与 +// pruneConfirmedUserPrefixTx 一致，使新授权的 observed baseline 和 retained floor 不会 +// 交叉提交成静默空洞。母行锁又能在首次 authorization 尚不存在时串行化同一 +// raw auth key 的并发登录/换号。 +func bindAuthorization(ctx context.Context, db sqlcgen.DBTX, a domain.Authorization) error { + keyID := authKeyIDToInt64(a.AuthKeyID) + var lockedKeyID int64 + if err := db.QueryRow(ctx, ` +SELECT auth_key_id +FROM auth_keys +WHERE auth_key_id = $1 +FOR UPDATE`, keyID).Scan(&lockedKeyID); err != nil { + return fmt.Errorf("lock auth key for authorization: %w", err) + } + + if _, err := db.Exec(ctx, ` +INSERT INTO user_update_watermarks (user_id, contiguous_pts) +VALUES ($1, 0) +ON CONFLICT (user_id) DO NOTHING`, a.UserID); err != nil { + return fmt.Errorf("ensure authorization user update watermark: %w", err) + } + var currentPts int + if err := db.QueryRow(ctx, ` +SELECT contiguous_pts +FROM user_update_watermarks +WHERE user_id = $1 +FOR UPDATE`, a.UserID).Scan(¤tPts); err != nil { + return fmt.Errorf("lock authorization user update watermark: %w", err) + } + if _, err := db.Exec(ctx, ` +INSERT INTO user_update_retention (user_id) +VALUES ($1) +ON CONFLICT (user_id) DO NOTHING`, a.UserID); err != nil { + return fmt.Errorf("ensure authorization user update retention: %w", err) + } + var retainedFloor int + if err := db.QueryRow(ctx, ` +SELECT retained_through_pts +FROM user_update_retention +WHERE user_id = $1 +FOR UPDATE`, a.UserID).Scan(&retainedFloor); err != nil { + return fmt.Errorf("lock authorization user update retention: %w", err) + } + if retainedFloor > currentPts { + return fmt.Errorf( + "authorization update baseline invariant violation: user %d retained floor %d exceeds contiguous watermark %d", + a.UserID, retainedFloor, currentPts, + ) + } + + // 每次 Bind 都是一次显式登录 baseline：delivered pts 推进到已锁定的账号连续水位； + // observed 只推进到已删除的 retained floor，不把 live tail 伪装成客户端确认。 + // 历史遗留的 state 若超出账号 contiguous watermark，必须 fail-fast；不得用 + // GREATEST 把非法 future cursor 保留下来。WHERE 也封住“预检后并发插入”的竞态。 + tag, err := db.Exec(ctx, ` +INSERT INTO update_states (auth_key_id, user_id, pts, qts, date, seq, observed_pts) +VALUES ($1, $2, $3, 0, EXTRACT(EPOCH FROM now())::int, 0, $4) +ON CONFLICT (auth_key_id, user_id) DO UPDATE SET + pts = GREATEST(update_states.pts, EXCLUDED.pts), + qts = GREATEST(update_states.qts, EXCLUDED.qts), + date = GREATEST(update_states.date, EXCLUDED.date), + seq = GREATEST(update_states.seq, EXCLUDED.seq), + observed_pts = GREATEST(update_states.observed_pts, EXCLUDED.observed_pts), + updated_at = now() +WHERE update_states.pts >= 0 + AND update_states.pts <= $3 + AND update_states.observed_pts <= $3`, keyID, a.UserID, currentPts, retainedFloor) + if err != nil { + return fmt.Errorf("upsert authorization update baseline: %w", err) + } + if tag.RowsAffected() != 1 { + return fmt.Errorf( + "authorization update baseline invariant violation: auth key %x user %d has pts or observed_pts outside contiguous watermark %d", + a.AuthKeyID, a.UserID, currentPts, + ) + } + + if _, err := db.Exec(ctx, ` +DELETE FROM update_states +WHERE auth_key_id = $1 + AND user_id <> $2`, keyID, a.UserID); err != nil { + return fmt.Errorf("delete stale cross-user update states: %w", err) + } + + if _, err := db.Exec(ctx, ` INSERT INTO authorizations (auth_key_id, user_id, hash, layer, device_model, platform, system_version, api_id, app_version, ip, password_pending) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11) ON CONFLICT (auth_key_id) DO UPDATE SET @@ -43,10 +147,9 @@ ON CONFLICT (auth_key_id) DO UPDATE SET ip = EXCLUDED.ip, password_pending = EXCLUDED.password_pending, active_at = now()`, - authKeyIDToInt64(a.AuthKeyID), a.UserID, a.Hash, int32(a.Layer), a.DeviceModel, a.Platform, a.SystemVersion, int32(a.APIID), a.AppVersion, a.IP, a.PasswordPending, - ) - if err != nil { - return fmt.Errorf("upsert authorization: %w", err) + keyID, a.UserID, a.Hash, int32(a.Layer), a.DeviceModel, a.Platform, a.SystemVersion, int32(a.APIID), a.AppVersion, a.IP, a.PasswordPending, + ); err != nil { + return fmt.Errorf("write authorization: %w", err) } return nil } @@ -129,7 +232,8 @@ RETURNING auth_key_id, user_id, hash, layer, device_model, platform, system_vers } // RevokeByHash 删除协议 auth_key 作为远程踢设备的持久化事实入口。 -// authorizations/update_states 通过 FK cascade 删除；关联 temp auth key 显式删除，避免 raw temp key 重连。 +// authorizations 通过 FK cascade 删除；update_states 没有 auth_keys FK，必须显式清理； +// 关联 temp auth key 也显式删除，避免 raw temp key 重连。 func (s *AuthorizationStore) RevokeByHash(ctx context.Context, userID, hash int64) (domain.Authorization, bool, error) { row := s.db.QueryRow(ctx, ` WITH target AS MATERIALIZED ( @@ -144,12 +248,18 @@ WITH target AS MATERIALIZED ( WHERE perm_auth_key_id IN (SELECT auth_key_id FROM target) ) RETURNING auth_key_id +), deleted_update_states AS ( + DELETE FROM update_states + WHERE auth_key_id IN (SELECT auth_key_id FROM target) + RETURNING auth_key_id ), deleted_keys AS ( DELETE FROM auth_keys WHERE auth_key_id IN (SELECT auth_key_id FROM target) RETURNING auth_key_id ), touched AS ( - SELECT count(*) FROM deleted_temp + SELECT + (SELECT count(*) FROM deleted_temp) + + (SELECT count(*) FROM deleted_update_states) AS count ) SELECT target.auth_key_id, target.user_id, target.hash, target.layer, target.device_model, target.platform, target.system_version, target.api_id, target.app_version, target.ip, target.password_pending, @@ -207,12 +317,18 @@ WITH target AS MATERIALIZED ( WHERE perm_auth_key_id IN (SELECT auth_key_id FROM target) ) RETURNING auth_key_id +), deleted_update_states AS ( + DELETE FROM update_states + WHERE auth_key_id IN (SELECT auth_key_id FROM target) + RETURNING auth_key_id ), deleted_keys AS ( DELETE FROM auth_keys WHERE auth_key_id IN (SELECT auth_key_id FROM target) RETURNING auth_key_id ), touched AS ( - SELECT count(*) FROM deleted_temp + SELECT + (SELECT count(*) FROM deleted_temp) + + (SELECT count(*) FROM deleted_update_states) AS count ) SELECT target.auth_key_id, target.user_id, target.hash, target.layer, target.device_model, target.platform, target.system_version, target.api_id, target.app_version, target.ip, target.password_pending, diff --git a/internal/store/postgres/bootstrap_update_job.go b/internal/store/postgres/bootstrap_update_job.go index 173b44d8..a37c438d 100644 --- a/internal/store/postgres/bootstrap_update_job.go +++ b/internal/store/postgres/bootstrap_update_job.go @@ -60,11 +60,11 @@ func (s *BootstrapUpdateJobStore) MarkReadyForSession(ctx context.Context, userI tag, err := s.db.Exec(ctx, ` UPDATE bootstrap_update_jobs SET status = 'ready', + session_id = $3, ready_at = now(), updated_at = now() WHERE user_id = $1 AND auth_key_id = $2 - AND session_id = $3 AND status = 'pending'`, userID, authKeyIDToInt64(authKeyID), sessionID) if err != nil { diff --git a/internal/store/postgres/bootstrap_update_job_integration_test.go b/internal/store/postgres/bootstrap_update_job_integration_test.go new file mode 100644 index 00000000..7cffb304 --- /dev/null +++ b/internal/store/postgres/bootstrap_update_job_integration_test.go @@ -0,0 +1,53 @@ +package postgres + +import ( + "context" + "testing" + "time" + + "telesrv/internal/domain" +) + +func TestBootstrapUpdateJobPostgresSameAuthKeyReconnectTakesOverPendingSession(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + user := createLoginCodeDeliveryTestUser(t, ctx, pool, "bootstrap-reconnect") + msg, err := NewMessageStore(pool).Create(ctx, domain.Message{ + OwnerUserID: user.ID, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID}, + Date: int(time.Now().Unix()), + Body: "Login code: 12345", + }) + if err != nil { + t.Fatalf("create bootstrap message: %v", err) + } + bootstrap := NewBootstrapUpdateJobStore(pool) + authKeyID := [8]byte{1, 3, 5, 7} + const ( + oldSessionID = int64(11001) + newSessionID = int64(22002) + ) + job, err := bootstrap.EnqueueLoginMessage(ctx, domain.BootstrapUpdateJob{ + Kind: domain.BootstrapUpdateJobLoginMessage, UserID: user.ID, + AuthKeyID: authKeyID, SessionID: oldSessionID, MessageBoxID: msg.ID, + }) + if err != nil { + t.Fatalf("enqueue bootstrap: %v", err) + } + if ready, err := bootstrap.MarkReadyForSession(ctx, user.ID, [8]byte{9}, newSessionID); err != nil || ready != 0 { + t.Fatalf("different-auth ready=%d err=%v, want 0/nil", ready, err) + } + ready, err := bootstrap.MarkReadyForSession(ctx, user.ID, authKeyID, newSessionID) + if err != nil || ready != 1 { + t.Fatalf("same-auth reconnect ready=%d err=%v, want 1/nil", ready, err) + } + var status string + var sessionID int64 + if err := pool.QueryRow(ctx, `SELECT status, session_id FROM bootstrap_update_jobs WHERE id = $1`, job.ID).Scan(&status, &sessionID); err != nil { + t.Fatalf("load bootstrap job: %v", err) + } + if status != string(domain.BootstrapUpdateJobReady) || sessionID != newSessionID { + t.Fatalf("bootstrap status/session = %s/%d, want ready/%d", status, sessionID, newSessionID) + } +} diff --git a/internal/store/postgres/channel_difference_integration_test.go b/internal/store/postgres/channel_difference_integration_test.go index fbda74c7..842949e9 100644 --- a/internal/store/postgres/channel_difference_integration_test.go +++ b/internal/store/postgres/channel_difference_integration_test.go @@ -3,6 +3,7 @@ package postgres import ( "context" "telesrv/internal/domain" + "telesrv/internal/store" "testing" ) @@ -173,7 +174,6 @@ func TestChannelStorePublicPreviewDifferenceSkipsNonMemberMessages(t *testing.T) if err != nil { t.Fatalf("send channel message: %v", err) } - diff, err := channels.ListChannelDifference(ctx, domain.ChannelDifferenceRequest{ UserID: viewer.ID, ChannelID: channelID, @@ -233,16 +233,37 @@ func TestChannelStoreDifferenceUsesDurableMessageSnapshots(t *testing.T) { t.Fatalf("create channel: %v", err) } channelID = created.Channel.ID - sent, err := channels.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + sendReq := domain.SendChannelMessageRequest{ UserID: owner.ID, ChannelID: channelID, RandomID: 941, Message: "original", Date: 1700000381, - }) + } + sent, err := channels.SendChannelMessage(ctx, sendReq) if err != nil { t.Fatalf("send channel message: %v", err) } + fingerprint, err := store.ChannelSendFingerprint(sendReq) + if err != nil { + t.Fatalf("fingerprint channel message: %v", err) + } + replayReq := domain.ChannelSendReplayRequest{ChannelID: channelID, SenderUserID: owner.ID, RandomID: sent.Message.RandomID, IdempotencyFingerprint: fingerprint} + type replayState struct { + pts int + events int + } + loadReplayState := func() replayState { + t.Helper() + var state replayState + if err := pool.QueryRow(ctx, `SELECT pts FROM channels WHERE id = $1`, channelID).Scan(&state.pts); err != nil { + t.Fatalf("load channel pts: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM channel_update_events WHERE channel_id = $1`, channelID).Scan(&state.events); err != nil { + t.Fatalf("count channel events: %v", err) + } + return state + } if _, err := channels.EditChannelMessage(ctx, domain.EditChannelMessageRequest{ UserID: owner.ID, ChannelID: channelID, @@ -261,12 +282,16 @@ func TestChannelStoreDifferenceUsesDurableMessageSnapshots(t *testing.T) { }); err != nil { t.Fatalf("second edit: %v", err) } - duplicate, found, err := channels.duplicateChannelMessage(ctx, channelID, owner.ID, sent.Message.RandomID) + beforeReplay := loadReplayState() + duplicate, found, err := channels.LookupChannelSendReplay(ctx, replayReq) if err != nil { t.Fatalf("duplicate channel message: %v", err) } - if !found || !duplicate.Duplicate || duplicate.Event.Type != domain.ChannelUpdateNewMessage || duplicate.Message.Body != "original" || duplicate.Event.Message.Body != "original" { - t.Fatalf("duplicate after edit = %+v found=%v, want original new-message snapshot", duplicate, found) + if !found || !duplicate.Duplicate || duplicate.Event.Type != domain.ChannelUpdateNewMessage || duplicate.Message.Body != "second edit" || duplicate.Event.Message.Body != "second edit" || duplicate.Event.Pts != sent.Event.Pts { + t.Fatalf("duplicate after edit = %+v found=%v, want current snapshot with first-send pts", duplicate, found) + } + if after := loadReplayState(); after != beforeReplay { + t.Fatalf("edit replay mutated channel state = %+v, want %+v", after, beforeReplay) } diff, err := channels.ListChannelDifference(ctx, domain.ChannelDifferenceRequest{ @@ -287,6 +312,26 @@ func TestChannelStoreDifferenceUsesDurableMessageSnapshots(t *testing.T) { if diff.OtherUpdates[0].Message.Body != "first edit" || diff.OtherUpdates[1].Message.Body != "second edit" { t.Fatalf("edit snapshots = %q/%q, want first edit/second edit", diff.OtherUpdates[0].Message.Body, diff.OtherUpdates[1].Message.Body) } + deleted, err := channels.DeleteChannelMessages(ctx, domain.DeleteChannelMessagesRequest{ + UserID: owner.ID, ChannelID: channelID, IDs: []int{sent.Message.ID}, Date: 1700000384, + }) + if err != nil { + t.Fatalf("delete channel message: %v", err) + } + beforeReplay = loadReplayState() + duplicate, found, err = channels.LookupChannelSendReplay(ctx, replayReq) + if err != nil { + t.Fatalf("duplicate deleted channel message: %v", err) + } + if !found || !duplicate.Duplicate || duplicate.Message.Body != "original" || duplicate.Message.Pts != sent.Message.Pts || duplicate.Event.Pts != sent.Event.Pts { + t.Fatalf("duplicate after delete = %+v found=%v, want immutable first-send snapshot", duplicate, found) + } + if duplicate.ReplayDeleteEvent == nil || duplicate.ReplayDeleteEvent.Pts != deleted.Event.Pts || len(duplicate.ReplayDeleteEvent.MessageIDs) != 1 || duplicate.ReplayDeleteEvent.MessageIDs[0] != sent.Message.ID { + t.Fatalf("duplicate delete receipt = %+v, want durable event %+v", duplicate.ReplayDeleteEvent, deleted.Event) + } + if after := loadReplayState(); after != beforeReplay { + t.Fatalf("delete replay mutated channel state = %+v, want %+v", after, beforeReplay) + } } func TestChannelStoreSendFailureBeforePtsAllocationDoesNotRecordNoopGap(t *testing.T) { @@ -505,6 +550,13 @@ func TestReserveChannelPtsRollsBackWithTransaction(t *testing.T) { if got != created.Channel.Pts { t.Fatalf("channel pts after rollback = %d, want unchanged %d", got, created.Channel.Pts) } + batch, err := channels.MaxChannelPtsBatch(ctx, []int64{channelID, -channelID, channelID}) + if err != nil { + t.Fatalf("MaxChannelPtsBatch: %v", err) + } + if len(batch) != 1 || batch[channelID] != created.Channel.Pts { + t.Fatalf("batch channel pts = %v, want only %d:%d", batch, channelID, created.Channel.Pts) + } } func TestChannelStoreDifferenceTooLongSnapshot(t *testing.T) { diff --git a/internal/store/postgres/channel_helpers.go b/internal/store/postgres/channel_helpers.go index 3237d032..e9efa277 100644 --- a/internal/store/postgres/channel_helpers.go +++ b/internal/store/postgres/channel_helpers.go @@ -332,17 +332,12 @@ func (s *ChannelStore) ListDirtyActiveChannelsForUser(ctx context.Context, userI SELECT i.channel_id, c.pts FROM user_channel_member_index i JOIN channels c ON c.id = i.channel_id AND NOT c.deleted +JOIN channel_update_checkpoints cp ON cp.channel_id = i.channel_id WHERE i.user_id = $1 AND i.status = 'active' AND NOT i.deleted AND i.channel_id > $3 - AND EXISTS ( - SELECT 1 - FROM channel_update_events e - WHERE e.channel_id = i.channel_id - AND e.date > $2 - LIMIT 1 - ) + AND cp.latest_event_date > $2 ORDER BY i.channel_id ASC LIMIT $4`, userID, sinceDate, afterChannelID, limit) if err != nil { diff --git a/internal/store/postgres/channel_message_delete.go b/internal/store/postgres/channel_message_delete.go index e1bff144..288ab52e 100644 --- a/internal/store/postgres/channel_message_delete.go +++ b/internal/store/postgres/channel_message_delete.go @@ -379,8 +379,14 @@ ORDER BY id`, channel.ID, id32) deleted32 := int32s(deleted) if _, err := tx.Exec(ctx, ` UPDATE channel_messages -SET deleted = true, pts = $3, updated_at = now() -WHERE channel_id = $1 AND id = ANY($2::int[])`, channel.ID, deleted32, pts); err != nil { +SET deleted = true, + pts = $3, + delete_pts = $3, + delete_pts_count = $4, + delete_date = $5, + delete_message_ids = to_jsonb($2::int[]), + updated_at = now() +WHERE channel_id = $1 AND id = ANY($2::int[])`, channel.ID, deleted32, pts, len(deleted), date); err != nil { return nil, domain.ChannelUpdateEvent{}, channel, fmt.Errorf("soft delete channel messages: %w", err) } if err := deleteChannelUnreadMentionsTx(ctx, tx, channel.ID, deleted); err != nil { diff --git a/internal/store/postgres/channel_message_send.go b/internal/store/postgres/channel_message_send.go index 8109a4f6..6cedece8 100644 --- a/internal/store/postgres/channel_message_send.go +++ b/internal/store/postgres/channel_message_send.go @@ -8,18 +8,26 @@ import ( "github.com/jackc/pgx/v5" "strings" "telesrv/internal/domain" + "telesrv/internal/store" ) func (s *ChannelStore) SendChannelMessage(ctx context.Context, req domain.SendChannelMessageRequest) (domain.SendChannelMessageResult, error) { if req.UserID == 0 || req.ChannelID == 0 || (strings.TrimSpace(req.Message) == "" && req.Action == nil && req.Media.IsZero() && req.RichMessage.IsZero()) { return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid } + requestFingerprint, err := store.ChannelSendFingerprint(req) + if err != nil { + return domain.SendChannelMessageResult{}, err + } + // Normalize the fallback to an explicit receipt so retries of the internal + // transaction use exactly the same bytes as the first attempt. + req.IdempotencyFingerprint = requestFingerprint if req.Date == 0 { req.Date = nowUnix() } var lastErr error for attempt := 0; attempt < retryableChannelTxAttempts; attempt++ { - res, err := s.sendChannelMessageOnce(ctx, req) + res, err := s.sendChannelMessageOnce(ctx, req, requestFingerprint) if err == nil || !isRetryablePostgresTxError(err) || ctx.Err() != nil { return res, err } @@ -28,9 +36,14 @@ func (s *ChannelStore) SendChannelMessage(ctx context.Context, req domain.SendCh return domain.SendChannelMessageResult{}, lastErr } -func (s *ChannelStore) sendChannelMessageOnce(ctx context.Context, req domain.SendChannelMessageRequest) (domain.SendChannelMessageResult, error) { - if req.RandomID != 0 { - if dup, found, err := s.duplicateChannelMessage(ctx, req.ChannelID, req.UserID, req.RandomID); err != nil { +func (s *ChannelStore) sendChannelMessageOnce(ctx context.Context, req domain.SendChannelMessageRequest, requestFingerprint []byte) (domain.SendChannelMessageResult, error) { + if req.RandomID != 0 && !req.IdempotencyPreflighted { + if dup, found, err := s.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{ + ChannelID: req.ChannelID, + SenderUserID: req.UserID, + RandomID: req.RandomID, + IdempotencyFingerprint: requestFingerprint, + }); err != nil { return domain.SendChannelMessageResult{}, err } else if found { return dup, nil @@ -207,13 +220,30 @@ func (s *ChannelStore) sendChannelMessageOnce(ctx context.Context, req domain.Se Message: msg, SenderUserID: req.UserID, } - if err := insertChannelMessageTx(ctx, tx, msg); err != nil { + if err := insertChannelMessageWithFingerprintTx(ctx, tx, msg, requestFingerprint); err != nil { if isUniqueViolation(err) { - dup, found, dupErr := s.duplicateChannelMessage(ctx, req.ChannelID, req.UserID, req.RandomID) - if dupErr != nil || !found { + if req.RandomID == 0 { + return domain.SendChannelMessageResult{}, err + } + // A failed statement leaves the transaction aborted while its pool + // connection remains checked out. Release it before the winner lookup; + // otherwise a one-connection pool deadlocks waiting on itself. + if rollbackErr := tx.Rollback(ctx); rollbackErr != nil && !errors.Is(rollbackErr, pgx.ErrTxClosed) { + return domain.SendChannelMessageResult{}, fmt.Errorf("rollback channel random_id conflict: %w", rollbackErr) + } + committed = true // transaction is finalized by rollback; suppress deferred rollback + dup, found, dupErr := s.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{ + ChannelID: req.ChannelID, + SenderUserID: req.UserID, + RandomID: req.RandomID, + IdempotencyFingerprint: requestFingerprint, + }) + if dupErr != nil { return domain.SendChannelMessageResult{}, dupErr } - dup.Duplicate = true + if !found { + return domain.SendChannelMessageResult{}, fmt.Errorf("channel random_id unique conflict without replay receipt") + } return dup, nil } return domain.SendChannelMessageResult{}, err @@ -320,8 +350,33 @@ func filterSkippedChannelRecipients(recipients []int64, skip map[int64]struct{}) return out } -func (s *ChannelStore) duplicateChannelMessage(ctx context.Context, channelID, userID, randomID int64) (domain.SendChannelMessageResult, bool, error) { - row := s.db.QueryRow(ctx, `SELECT `+channelMessageColumns+` FROM channel_messages WHERE channel_id = $1 AND sender_user_id = $2 AND random_id = $3`, channelID, userID, randomID) +// LookupChannelSendReplay reads an immutable random_id receipt without running +// membership, permission, slow-mode, source/media resolution or allocation. A +// zero SavedPeer selects an ordinary channel receipt; monoforum sub-dialogs are +// scoped by the complete saved peer. +func (s *ChannelStore) LookupChannelSendReplay(ctx context.Context, lookup domain.ChannelSendReplayRequest) (domain.SendChannelMessageResult, bool, error) { + if lookup.ChannelID == 0 || lookup.SenderUserID == 0 || lookup.RandomID == 0 { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("channel send replay: invalid scope") + } + if err := store.ValidateSendFingerprint(lookup.IdempotencyFingerprint, "channel send replay"); err != nil { + return domain.SendChannelMessageResult{}, false, err + } + var row pgx.Row + if lookup.SavedPeer.ID == 0 { + if lookup.SavedPeer.Type != "" { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("channel send replay: incomplete saved peer scope") + } + row = s.db.QueryRow(ctx, `SELECT `+channelMessageColumns+` FROM channel_messages +WHERE channel_id = $1 AND sender_user_id = $2 AND saved_peer_type = '' AND saved_peer_id = 0 AND random_id = $3`, + lookup.ChannelID, lookup.SenderUserID, lookup.RandomID) + } else { + if lookup.SavedPeer.Type != domain.PeerTypeUser { + return domain.SendChannelMessageResult{}, false, fmt.Errorf("channel send replay: invalid saved peer scope") + } + row = s.db.QueryRow(ctx, `SELECT `+channelMessageColumns+` FROM channel_messages +WHERE channel_id = $1 AND sender_user_id = $2 AND saved_peer_type = $3 AND saved_peer_id = $4 AND random_id = $5`, + lookup.ChannelID, lookup.SenderUserID, string(lookup.SavedPeer.Type), lookup.SavedPeer.ID, lookup.RandomID) + } msg, err := scanChannelMessage(row) if errors.Is(err, pgx.ErrNoRows) { return domain.SendChannelMessageResult{}, false, nil @@ -329,18 +384,71 @@ func (s *ChannelStore) duplicateChannelMessage(ctx context.Context, channelID, u if err != nil { return domain.SendChannelMessageResult{}, false, err } - channel, err := getChannelByID(ctx, s.db, channelID) + result, err := s.channelDuplicateReplayResult(ctx, msg, lookup.IdempotencyFingerprint) if err != nil { return domain.SendChannelMessageResult{}, false, err } - event, err := s.eventForChannelMessage(ctx, channelID, msg.ID) + result.Duplicate = true + return result, true, nil +} + +func (s *ChannelStore) channelDuplicateReplayResult(ctx context.Context, msg domain.ChannelMessage, expectedFingerprint []byte) (domain.SendChannelMessageResult, error) { + var storedFingerprint []byte + var snapshotJSON, deleteIDsJSON string + var deletePts, deletePtsCount, deleteDate int + if err := s.db.QueryRow(ctx, ` +SELECT request_fingerprint, send_snapshot::text, delete_pts, delete_pts_count, delete_date, delete_message_ids::text +FROM channel_messages +WHERE channel_id = $1 AND id = $2`, msg.ChannelID, msg.ID).Scan( + &storedFingerprint, &snapshotJSON, &deletePts, &deletePtsCount, &deleteDate, &deleteIDsJSON, + ); err != nil { + return domain.SendChannelMessageResult{}, err + } + if !store.SameSendFingerprint(storedFingerprint, expectedFingerprint) { + return domain.SendChannelMessageResult{}, domain.ErrMessageRandomIDDuplicate + } + first, err := store.DecodeChannelSendSnapshot([]byte(snapshotJSON)) if err != nil { - return domain.SendChannelMessageResult{}, false, err + return domain.SendChannelMessageResult{}, fmt.Errorf("decode duplicate channel message %d snapshot: %w", msg.ID, err) } - if event.Message.ID != 0 { - msg = event.Message + if first.ChannelID != msg.ChannelID || first.ID != msg.ID || first.SenderUserID != msg.SenderUserID || first.RandomID != msg.RandomID || first.SavedPeer != msg.SavedPeer { + return domain.SendChannelMessageResult{}, fmt.Errorf("duplicate channel message %d snapshot disagrees with random_id receipt", msg.ID) } - return domain.SendChannelMessageResult{Channel: channel, Message: msg, Event: event, Duplicate: true}, true, nil + channel, err := getChannelByID(ctx, s.db, msg.ChannelID) + if err != nil { + return domain.SendChannelMessageResult{}, err + } + replay := msg + var replayDelete *domain.ChannelUpdateEvent + if msg.Deleted { + replay = first + messageIDs, err := decodeEventMessageIDs(deleteIDsJSON) + if err != nil { + return domain.SendChannelMessageResult{}, fmt.Errorf("decode duplicate channel message %d delete ids: %w", msg.ID, err) + } + if deletePts <= 0 || deletePtsCount <= 0 || len(messageIDs) == 0 { + return domain.SendChannelMessageResult{}, fmt.Errorf("duplicate channel message %d is deleted without a durable delete receipt", msg.ID) + } + deleteEvent := domain.ChannelUpdateEvent{ + ChannelID: msg.ChannelID, + Type: domain.ChannelUpdateDeleteMessages, + Pts: deletePts, + PtsCount: deletePtsCount, + Date: deleteDate, + MessageIDs: messageIDs, + } + replayDelete = &deleteEvent + } + event := domain.ChannelUpdateEvent{ + ChannelID: msg.ChannelID, + Type: domain.ChannelUpdateNewMessage, + Pts: first.Pts, + PtsCount: 1, + Date: first.Date, + Message: replay, + SenderUserID: first.SenderUserID, + } + return domain.SendChannelMessageResult{Channel: channel, Message: replay, Event: event, Duplicate: true, ReplayDeleteEvent: replayDelete}, nil } func (s *ChannelStore) insertServiceMessage(ctx context.Context, tx pgx.Tx, channel domain.Channel, senderUserID int64, date int, action domain.ChannelMessageAction) (domain.ChannelMessage, domain.ChannelUpdateEvent, error) { @@ -400,6 +508,26 @@ func channelServiceActionForMessage(channelID int64, msgID int, action domain.Ch } func insertChannelMessageTx(ctx context.Context, tx pgx.Tx, msg domain.ChannelMessage) error { + return insertChannelMessageWithFingerprintTx(ctx, tx, msg, nil) +} + +// insertChannelMessageWithFingerprintTx is the only first-send write boundary +// for client-random-id channel messages. Callers that create service/discussion +// rows without random_id use insertChannelMessageTx and persist the legacy-safe +// empty default instead. +func insertChannelMessageWithFingerprintTx(ctx context.Context, tx pgx.Tx, msg domain.ChannelMessage, requestFingerprint []byte) error { + if msg.RandomID != 0 { + if err := store.ValidateSendFingerprint(requestFingerprint, "insert channel message"); err != nil { + return err + } + } else if len(requestFingerprint) != 0 { + if err := store.ValidateSendFingerprint(requestFingerprint, "insert channel message"); err != nil { + return err + } + } + if requestFingerprint == nil { + requestFingerprint = []byte{} + } entities, err := encodeMessageEntities(msg.Entities) if err != nil { return err @@ -428,6 +556,13 @@ func insertChannelMessageTx(ctx context.Context, tx pgx.Tx, msg domain.ChannelMe if err != nil { return err } + sendSnapshot := []byte("{}") + if msg.RandomID != 0 { + sendSnapshot, err = store.EncodeChannelSendSnapshot(msg) + if err != nil { + return err + } + } var sendAsType sql.NullString var sendAsID sql.NullInt64 if msg.SendAs != nil && msg.SendAs.ID != 0 { @@ -456,12 +591,12 @@ INSERT INTO channel_messages ( channel_id, id, random_id, sender_user_id, from_peer_type, from_peer_id, send_as_peer_type, send_as_peer_id, message_date, edit_date, post, silent, noforwards, body, entities, reply_to, reply_to_msg_id, reply_to_peer_type, reply_to_peer_id, reply_to_top_id, - fwd_from, discussion_channel_id, discussion_message_id, action, pts, deleted, media, reply_markup, rich_message, ttl_period, expires_at, post_author, via_bot_id, from_boosts_applied, grouped_id, saved_peer_type, saved_peer_id -) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23,$24,$25,$26,$27,$28,$29,$30,$31,$32,$33,$34,$35,$36,$37)`, + fwd_from, discussion_channel_id, discussion_message_id, action, pts, deleted, media, reply_markup, rich_message, ttl_period, expires_at, post_author, via_bot_id, from_boosts_applied, grouped_id, saved_peer_type, saved_peer_id, send_snapshot, request_fingerprint +) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23,$24,$25,$26,$27,$28,$29,$30,$31,$32,$33,$34,$35,$36,$37,$38::jsonb,$39::bytea)`, msg.ChannelID, msg.ID, msg.RandomID, msg.SenderUserID, string(msg.From.Type), msg.From.ID, sendAsType, sendAsID, msg.Date, msg.EditDate, msg.Post, msg.Silent, msg.NoForwards, msg.Body, entities, reply, replyMsgID, replyPeerType, replyPeerID, replyTopID, - forward, discussionChannelID, discussionMessageID, action, msg.Pts, msg.Deleted, media, replyMarkup, richMessage, msg.TTLPeriod, msg.ExpiresAt, msg.PostAuthor, msg.ViaBotID, msg.FromBoostsApplied, msg.GroupedID, string(msg.SavedPeer.Type), msg.SavedPeer.ID); err != nil { + forward, discussionChannelID, discussionMessageID, action, msg.Pts, msg.Deleted, media, replyMarkup, richMessage, msg.TTLPeriod, msg.ExpiresAt, msg.PostAuthor, msg.ViaBotID, msg.FromBoostsApplied, msg.GroupedID, string(msg.SavedPeer.Type), msg.SavedPeer.ID, sendSnapshot, requestFingerprint); err != nil { return fmt.Errorf("insert channel message: %w", err) } // 共享媒体索引(迁移 0118):创建即按媒体类别建索引行,供 messages.search 媒体标签页。 diff --git a/internal/store/postgres/channel_monoforum.go b/internal/store/postgres/channel_monoforum.go index f47397ee..3e8fd0a9 100644 --- a/internal/store/postgres/channel_monoforum.go +++ b/internal/store/postgres/channel_monoforum.go @@ -9,6 +9,7 @@ import ( "github.com/jackc/pgx/v5" "telesrv/internal/domain" + "telesrv/internal/store" ) // SendMonoforumMessage 向 monoforum(频道私信)虚拟频道发一条消息,按 saved_peer 分订阅者子会话。 @@ -19,11 +20,22 @@ func (s *ChannelStore) SendMonoforumMessage(ctx context.Context, req domain.Send req.SavedPeer.Type != domain.PeerTypeUser || strings.TrimSpace(req.Message) == "" { return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid } + requestFingerprint, err := store.MonoforumSendFingerprint(req) + if err != nil { + return domain.SendChannelMessageResult{}, err + } + req.IdempotencyFingerprint = requestFingerprint if req.Date == 0 { req.Date = nowUnix() } - if req.RandomID != 0 { - if dup, found, err := s.duplicateMonoforumMessage(ctx, req.MonoforumID, req.SenderUserID, req.SavedPeer, req.RandomID); err != nil { + if req.RandomID != 0 && !req.IdempotencyPreflighted { + if dup, found, err := s.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{ + ChannelID: req.MonoforumID, + SenderUserID: req.SenderUserID, + SavedPeer: req.SavedPeer, + RandomID: req.RandomID, + IdempotencyFingerprint: requestFingerprint, + }); err != nil { return domain.SendChannelMessageResult{}, err } else if found { return dup, nil @@ -82,15 +94,32 @@ func (s *ChannelStore) SendMonoforumMessage(ctx context.Context, req domain.Send Message: msg, SenderUserID: req.SenderUserID, } - if err := insertChannelMessageTx(ctx, tx, msg); err != nil { + if err := insertChannelMessageWithFingerprintTx(ctx, tx, msg, requestFingerprint); err != nil { if isUniqueViolation(err) { - // 唯一约束按 (channel,sender,random_id) 三元组;只有同一订阅者子会话的真重发才算重复。 - // 跨子会话复用同一 random_id(异常客户端)按 saved_peer 过滤后命中不到 → 干净返错,不串消息。 - dup, found, dupErr := s.duplicateMonoforumMessage(ctx, req.MonoforumID, req.SenderUserID, req.SavedPeer, req.RandomID) - if dupErr != nil || !found { + if req.RandomID == 0 { + return domain.SendChannelMessageResult{}, err + } + // The winner lookup must not ask the pool for a second connection + // while this aborted transaction still owns the first one. + if rollbackErr := tx.Rollback(ctx); rollbackErr != nil && !errors.Is(rollbackErr, pgx.ErrTxClosed) { + return domain.SendChannelMessageResult{}, fmt.Errorf("rollback monoforum random_id conflict: %w", rollbackErr) + } + committed = true // transaction is finalized by rollback; suppress deferred rollback + // The four-column unique scope is only the race fence. Acceptance + // still requires the exact immutable request fingerprint. + dup, found, dupErr := s.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{ + ChannelID: req.MonoforumID, + SenderUserID: req.SenderUserID, + SavedPeer: req.SavedPeer, + RandomID: req.RandomID, + IdempotencyFingerprint: requestFingerprint, + }) + if dupErr != nil { return domain.SendChannelMessageResult{}, dupErr } - dup.Duplicate = true + if !found { + return domain.SendChannelMessageResult{}, fmt.Errorf("monoforum random_id unique conflict without replay receipt") + } return dup, nil } return domain.SendChannelMessageResult{}, err @@ -183,33 +212,6 @@ func (s *ChannelStore) ResolveMonoforumSend(ctx context.Context, viewerUserID, m return mono, isAdmin, nil } -// duplicateMonoforumMessage 按 (channel,sender,saved_peer,random_id) 查重发,确保同一发件人向不同 -// 订阅者子会话用相同 random_id 时不会互相误判为重复。 -func (s *ChannelStore) duplicateMonoforumMessage(ctx context.Context, channelID, senderUserID int64, savedPeer domain.Peer, randomID int64) (domain.SendChannelMessageResult, bool, error) { - row := s.db.QueryRow(ctx, `SELECT `+channelMessageColumns+` FROM channel_messages -WHERE channel_id = $1 AND sender_user_id = $2 AND saved_peer_type = $3 AND saved_peer_id = $4 AND random_id = $5`, - channelID, senderUserID, string(savedPeer.Type), savedPeer.ID, randomID) - msg, err := scanChannelMessage(row) - if errors.Is(err, pgx.ErrNoRows) { - return domain.SendChannelMessageResult{}, false, nil - } - if err != nil { - return domain.SendChannelMessageResult{}, false, err - } - channel, err := getChannelByID(ctx, s.db, channelID) - if err != nil { - return domain.SendChannelMessageResult{}, false, err - } - event, err := s.eventForChannelMessage(ctx, channelID, msg.ID) - if err != nil { - return domain.SendChannelMessageResult{}, false, err - } - if event.Message.ID != 0 { - msg = event.Message - } - return domain.SendChannelMessageResult{Channel: channel, Message: msg, Event: event, Duplicate: true}, true, nil -} - // ListMonoforumDialogs 列出 monoforum 的订阅者子会话(每个 saved_peer 一条,取其 top 消息), // 按 top 消息 id 倒序分页。走部分索引 channel_messages_monoforum_sublist_idx。 func (s *ChannelStore) ListMonoforumDialogs(ctx context.Context, filter domain.MonoforumDialogsFilter) (domain.MonoforumDialogList, error) { diff --git a/internal/store/postgres/channel_monoforum_send_integration_test.go b/internal/store/postgres/channel_monoforum_send_integration_test.go index edbc0928..64bc5c1b 100644 --- a/internal/store/postgres/channel_monoforum_send_integration_test.go +++ b/internal/store/postgres/channel_monoforum_send_integration_test.go @@ -2,9 +2,11 @@ package postgres import ( "context" + "errors" "testing" "telesrv/internal/domain" + "telesrv/internal/store" ) // TestSendMonoforumMessageAndHistoryPostgres 回归频道私信(monoforum)发送+读历史的 PG 实现: @@ -95,6 +97,16 @@ func TestSendMonoforumMessageAndHistoryPostgres(t *testing.T) { if !dup.Duplicate || dup.Message.ID != m1.Message.ID { t.Fatalf("dup = %+v, want duplicate of m1 id %d", dup.Message, m1.Message.ID) } + if _, err := channels.SendMonoforumMessage(ctx, domain.SendMonoforumMessageRequest{MonoforumID: monoID, SenderUserID: sub.ID, SavedPeer: subPeer, RandomID: 111, Message: "changed", Date: 1700001004}); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("changed monoforum intent err = %v, want ErrMessageRandomIDDuplicate", err) + } + var monoforumFingerprint []byte + if err := pool.QueryRow(ctx, `SELECT request_fingerprint FROM channel_messages WHERE channel_id=$1 AND id=$2`, monoID, m1.Message.ID).Scan(&monoforumFingerprint); err != nil { + t.Fatalf("load monoforum fingerprint: %v", err) + } + if len(monoforumFingerprint) != 32 { + t.Fatalf("monoforum fingerprint length = %d, want 32", len(monoforumFingerprint)) + } // 历史(经 scanChannelMessage 读回 saved_peer)。 hist, err := channels.ListMonoforumHistory(ctx, domain.MonoforumHistoryFilter{MonoforumID: monoID, SavedPeer: subPeer, Limit: 10}) @@ -136,6 +148,20 @@ func TestSendMonoforumMessageAndHistoryPostgres(t *testing.T) { if b.Duplicate || b.Message.ID == a.Message.ID { t.Fatalf("cross-sublist same random_id wrongly deduped: a=%d b=%d dup=%v", a.Message.ID, b.Message.ID, b.Duplicate) } + // SavedPeer belongs both to the lookup scope and to the fallback intent. + // Cross-sublist sends remain legal, while presenting another sublist's + // fingerprint for an existing scope must be rejected rather than replayed. + otherIntentFingerprint, err := store.MonoforumSendFingerprint(domain.SendMonoforumMessageRequest{ + MonoforumID: monoID, SenderUserID: owner.ID, SavedPeer: otherPeer, RandomID: 9001, Message: "to sub", + }) + if err != nil { + t.Fatalf("fingerprint mismatched saved peer: %v", err) + } + if _, _, err := channels.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{ + ChannelID: monoID, SenderUserID: owner.ID, SavedPeer: subPeer, RandomID: 9001, IdempotencyFingerprint: otherIntentFingerprint, + }); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("mismatched saved-peer fingerprint err = %v, want ErrMessageRandomIDDuplicate", err) + } // 同一子会话真重发仍去重。 again, err := channels.SendMonoforumMessage(ctx, domain.SendMonoforumMessageRequest{MonoforumID: monoID, SenderUserID: owner.ID, SavedPeer: subPeer, RandomID: 9001, Message: "to sub", Date: 1700001012}) if err != nil { @@ -159,4 +185,45 @@ func TestSendMonoforumMessageAndHistoryPostgres(t *testing.T) { if dialogs.Dialogs[1].SavedPeer != subPeer || dialogs.Dialogs[1].TopMessageID == 0 { t.Fatalf("dialogs[1] = %+v, want sub with top message", dialogs.Dialogs[1]) } + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin monoforum delete: %v", err) + } + mono, err := getChannelByID(ctx, tx, monoID) + if err != nil { + _ = tx.Rollback(ctx) + t.Fatalf("load monoforum for delete: %v", err) + } + _, deleteEvent, _, err := channels.deleteChannelMessagesTx(ctx, tx, mono, domain.ChannelMember{ChannelID: monoID, UserID: owner.ID, Role: domain.ChannelRoleCreator, Status: domain.ChannelMemberActive}, []int{a.Message.ID}, owner.ID, 1700001013) + if err != nil { + _ = tx.Rollback(ctx) + t.Fatalf("delete monoforum message: %v", err) + } + if err := tx.Commit(ctx); err != nil { + t.Fatalf("commit monoforum delete: %v", err) + } + var ptsBeforeReplay, eventsBeforeReplay int + if err := pool.QueryRow(ctx, `SELECT pts FROM channels WHERE id = $1`, monoID).Scan(&ptsBeforeReplay); err != nil { + t.Fatalf("load monoforum pts: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM channel_update_events WHERE channel_id = $1`, monoID).Scan(&eventsBeforeReplay); err != nil { + t.Fatalf("count monoforum events: %v", err) + } + deletedReplay, err := channels.SendMonoforumMessage(ctx, domain.SendMonoforumMessageRequest{MonoforumID: monoID, SenderUserID: owner.ID, SavedPeer: subPeer, RandomID: 9001, Message: "to sub", Date: 1700001014}) + if err != nil { + t.Fatalf("replay deleted monoforum message: %v", err) + } + if !deletedReplay.Duplicate || deletedReplay.Message.ID != a.Message.ID || deletedReplay.Message.Body != "to sub" || deletedReplay.ReplayDeleteEvent == nil || deletedReplay.ReplayDeleteEvent.Pts != deleteEvent.Pts { + t.Fatalf("deleted monoforum replay = %+v, want first snapshot + durable delete %+v", deletedReplay, deleteEvent) + } + var ptsAfterReplay, eventsAfterReplay int + if err := pool.QueryRow(ctx, `SELECT pts FROM channels WHERE id = $1`, monoID).Scan(&ptsAfterReplay); err != nil { + t.Fatalf("reload monoforum pts: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM channel_update_events WHERE channel_id = $1`, monoID).Scan(&eventsAfterReplay); err != nil { + t.Fatalf("recount monoforum events: %v", err) + } + if ptsAfterReplay != ptsBeforeReplay || eventsAfterReplay != eventsBeforeReplay { + t.Fatalf("deleted monoforum replay mutated pts/events = %d/%d, want %d/%d", ptsAfterReplay, eventsAfterReplay, ptsBeforeReplay, eventsBeforeReplay) + } } diff --git a/internal/store/postgres/channel_send_idempotency_integration_test.go b/internal/store/postgres/channel_send_idempotency_integration_test.go new file mode 100644 index 00000000..b0311ead --- /dev/null +++ b/internal/store/postgres/channel_send_idempotency_integration_test.go @@ -0,0 +1,490 @@ +package postgres + +import ( + "bytes" + "context" + "crypto/sha256" + "errors" + "os" + "strings" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + + "telesrv/deploy" + "telesrv/internal/domain" + "telesrv/internal/store" +) + +func TestChannelSendFingerprintReplayPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + + users := NewUserStore(pool) + owner, err := users.Create(ctx, domain.User{AccessHash: 181, Phone: "+1781" + suffix + "01", FirstName: "ChannelReplayOwner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + var channelID int64 + t.Cleanup(func() { + if channelID != 0 { + _, _ = pool.Exec(ctx, `DELETE FROM channels WHERE id = $1`, channelID) + } + _, _ = pool.Exec(ctx, `DELETE FROM users WHERE id = $1`, owner.ID) + }) + + channels := NewChannelStore(pool) + created, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, + Title: "Channel replay " + suffix, + Megagroup: true, + Date: 1700100000, + }) + if err != nil { + t.Fatalf("create channel: %v", err) + } + channelID = created.Channel.ID + base := domain.SendChannelMessageRequest{ + UserID: owner.ID, + ChannelID: channelID, + RandomID: 781001, + Message: "immutable original", + Entities: []domain.MessageEntity{{ + Type: domain.MessageEntityBold, + Offset: 0, + Length: 9, + }}, + Date: 1700100001, + } + wantFingerprint, err := store.ChannelSendFingerprint(base) + if err != nil { + t.Fatalf("fingerprint base: %v", err) + } + first, err := channels.SendChannelMessage(ctx, base) + if err != nil { + t.Fatalf("first send: %v", err) + } + var storedFingerprint []byte + if err := pool.QueryRow(ctx, `SELECT request_fingerprint FROM channel_messages WHERE channel_id = $1 AND id = $2`, channelID, first.Message.ID).Scan(&storedFingerprint); err != nil { + t.Fatalf("load fingerprint: %v", err) + } + if !bytes.Equal(storedFingerprint, wantFingerprint) { + t.Fatalf("stored fingerprint = %x, want %x", storedFingerprint, wantFingerprint) + } + + type durableState struct { + pts int + events int + rows int + } + loadState := func(randomID int64) durableState { + t.Helper() + var state durableState + if err := pool.QueryRow(ctx, `SELECT pts FROM channels WHERE id = $1`, channelID).Scan(&state.pts); err != nil { + t.Fatalf("load channel pts: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM channel_update_events WHERE channel_id = $1`, channelID).Scan(&state.events); err != nil { + t.Fatalf("count channel events: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM channel_messages WHERE channel_id = $1 AND sender_user_id = $2 AND random_id = $3`, channelID, owner.ID, randomID).Scan(&state.rows); err != nil { + t.Fatalf("count random receipt: %v", err) + } + return state + } + + before := loadState(base.RandomID) + exact := base + exact.Date += 100 // execution time is not part of immutable intent + replay, err := channels.SendChannelMessage(ctx, exact) + if err != nil { + t.Fatalf("exact replay: %v", err) + } + if !replay.Duplicate || replay.Message.ID != first.Message.ID || replay.Event.Pts != first.Event.Pts { + t.Fatalf("exact replay = %+v, want first id=%d pts=%d", replay, first.Message.ID, first.Event.Pts) + } + if after := loadState(base.RandomID); after != before { + t.Fatalf("exact replay mutated state = %+v, want %+v", after, before) + } + + conflicts := []struct { + name string + mutate func(*domain.SendChannelMessageRequest) + }{ + {name: "body", mutate: func(req *domain.SendChannelMessageRequest) { req.Message = "changed body" }}, + {name: "media", mutate: func(req *domain.SendChannelMessageRequest) { + req.Media = &domain.MessageMedia{Kind: domain.MessageMediaKindPhoto, Photo: &domain.Photo{ID: 781, AccessHash: 782, DCID: 2}} + }}, + {name: "reply", mutate: func(req *domain.SendChannelMessageRequest) { + req.ReplyTo = &domain.MessageReply{MessageID: first.Message.ID} + }}, + {name: "group", mutate: func(req *domain.SendChannelMessageRequest) { req.GroupedID = 781003 }}, + } + for _, tc := range conflicts { + t.Run("conflict_"+tc.name, func(t *testing.T) { + changed := base + tc.mutate(&changed) + if _, err := channels.SendChannelMessage(ctx, changed); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("changed %s replay err = %v, want ErrMessageRandomIDDuplicate", tc.name, err) + } + if after := loadState(base.RandomID); after != before { + t.Fatalf("changed %s replay mutated state = %+v, want %+v", tc.name, after, before) + } + }) + } + + if _, err := channels.EditChannelMessage(ctx, domain.EditChannelMessageRequest{ + UserID: owner.ID, ChannelID: channelID, ID: first.Message.ID, Message: "edited current", EditDate: 1700100002, + }); err != nil { + t.Fatalf("edit channel message: %v", err) + } + editedReplay, err := channels.SendChannelMessage(ctx, exact) + if err != nil { + t.Fatalf("replay edited message: %v", err) + } + if !editedReplay.Duplicate || editedReplay.Message.Body != "edited current" || editedReplay.Event.Pts != first.Event.Pts { + t.Fatalf("edited replay = %+v, want current projection with first pts", editedReplay) + } + deleted, err := channels.DeleteChannelMessages(ctx, domain.DeleteChannelMessagesRequest{ + UserID: owner.ID, ChannelID: channelID, IDs: []int{first.Message.ID}, Date: 1700100003, + }) + if err != nil { + t.Fatalf("delete channel message: %v", err) + } + deletedReplay, err := channels.SendChannelMessage(ctx, exact) + if err != nil { + t.Fatalf("replay deleted message: %v", err) + } + if !deletedReplay.Duplicate || deletedReplay.Message.Body != base.Message || deletedReplay.Message.ID != first.Message.ID || deletedReplay.ReplayDeleteEvent == nil || deletedReplay.ReplayDeleteEvent.Pts != deleted.Event.Pts { + t.Fatalf("deleted replay = %+v, want immutable first snapshot + delete receipt %+v", deletedReplay, deleted.Event) + } + + // A raw request-boundary fingerprint must be stored byte-for-byte rather + // than replaced with the domain fallback. + raw := sha256.Sum256([]byte("raw channel TL intent")) + rawReq := domain.SendChannelMessageRequest{ + UserID: owner.ID, ChannelID: channelID, RandomID: 781002, Message: "raw fingerprint", Date: 1700100010, + IdempotencyFingerprint: raw[:], + } + rawSent, err := channels.SendChannelMessage(ctx, rawReq) + if err != nil { + t.Fatalf("raw fingerprint send: %v", err) + } + storedFingerprint = nil + if err := pool.QueryRow(ctx, `SELECT request_fingerprint FROM channel_messages WHERE channel_id = $1 AND id = $2`, channelID, rawSent.Message.ID).Scan(&storedFingerprint); err != nil { + t.Fatalf("load raw fingerprint: %v", err) + } + if !bytes.Equal(storedFingerprint, raw[:]) { + t.Fatalf("stored raw fingerprint = %x, want %x", storedFingerprint, raw) + } + + // Simulate a rolling old writer that omits the new column. The empty + // default keeps the write compatible, but it is never accepted as replay. + legacyID := rawSent.Message.ID + 100 + legacyRandomID := int64(781099) + if _, err := pool.Exec(ctx, ` +INSERT INTO channel_messages (channel_id, id, random_id, sender_user_id, from_peer_id, message_date, pts, body) +VALUES ($1,$2,$3,$4,$4,$5,$6,$7)`, channelID, legacyID, legacyRandomID, owner.ID, 1700100020, rawSent.Event.Pts+100, "legacy unknown intent"); err != nil { + t.Fatalf("old-writer insert without fingerprint: %v", err) + } + var legacyFingerprint []byte + if err := pool.QueryRow(ctx, `SELECT request_fingerprint FROM channel_messages WHERE channel_id=$1 AND id=$2`, channelID, legacyID).Scan(&legacyFingerprint); err != nil { + t.Fatalf("load legacy fingerprint: %v", err) + } + if len(legacyFingerprint) != 0 { + t.Fatalf("legacy fingerprint length = %d, want empty", len(legacyFingerprint)) + } + legacyReq := domain.SendChannelMessageRequest{UserID: owner.ID, ChannelID: channelID, RandomID: legacyRandomID, Message: "legacy unknown intent", Date: 1700100021} + legacyExpected, err := store.ChannelSendFingerprint(legacyReq) + if err != nil { + t.Fatalf("fingerprint legacy retry: %v", err) + } + if _, _, err := channels.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{ + ChannelID: channelID, SenderUserID: owner.ID, RandomID: legacyRandomID, IdempotencyFingerprint: legacyExpected, + }); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("legacy empty lookup err = %v, want ErrMessageRandomIDDuplicate", err) + } + if _, err := channels.SendChannelMessage(ctx, legacyReq); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("legacy empty send err = %v, want ErrMessageRandomIDDuplicate", err) + } +} + +func TestChannelSendFingerprintConcurrentRacePostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + owner, err := users.Create(ctx, domain.User{AccessHash: 191, Phone: "+1781" + suffix + "11", FirstName: "ChannelRaceOwner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + var channelIDs []int64 + t.Cleanup(func() { + if len(channelIDs) != 0 { + _, _ = pool.Exec(ctx, `DELETE FROM channels WHERE id = ANY($1::bigint[])`, channelIDs) + } + _, _ = pool.Exec(ctx, `DELETE FROM users WHERE id = $1`, owner.ID) + }) + newChannel := func(title string) int64 { + t.Helper() + created, err := NewChannelStore(pool).CreateChannel(ctx, domain.CreateChannelRequest{CreatorUserID: owner.ID, Title: title + suffix, Megagroup: true, Date: 1700110000}) + if err != nil { + t.Fatalf("create %s channel: %v", title, err) + } + channelIDs = append(channelIDs, created.Channel.ID) + return created.Channel.ID + } + + run := func(reqs [2]domain.SendChannelMessageRequest) ([2]domain.SendChannelMessageResult, [2]error) { + t.Helper() + var results [2]domain.SendChannelMessageResult + var errs [2]error + start := make(chan struct{}) + var wg sync.WaitGroup + for i := range reqs { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + results[i], errs[i] = NewChannelStore(pool).SendChannelMessage(ctx, reqs[i]) + }(i) + } + close(start) + wg.Wait() + return results, errs + } + + exactChannelID := newChannel("exact race ") + exactReq := domain.SendChannelMessageRequest{ + UserID: owner.ID, ChannelID: exactChannelID, RandomID: 791001, Message: "same intent", Date: 1700110001, + IdempotencyPreflighted: true, + } + exactResults, exactErrs := run([2]domain.SendChannelMessageRequest{exactReq, exactReq}) + for i, err := range exactErrs { + if err != nil { + t.Fatalf("exact race result[%d] err = %v", i, err) + } + } + if exactResults[0].Message.ID != exactResults[1].Message.ID || exactResults[0].Duplicate == exactResults[1].Duplicate { + t.Fatalf("exact race results = %+v / %+v, want same id and one duplicate", exactResults[0], exactResults[1]) + } + assertChannelRandomReceiptCount(t, ctx, pool, exactChannelID, owner.ID, exactReq.RandomID, 1) + + conflictChannelID := newChannel("conflict race ") + conflictA := domain.SendChannelMessageRequest{ + UserID: owner.ID, ChannelID: conflictChannelID, RandomID: 791002, Message: "intent A", Date: 1700110010, + IdempotencyPreflighted: true, + } + conflictB := conflictA + conflictB.Message = "intent B" + conflictResults, conflictErrs := run([2]domain.SendChannelMessageRequest{conflictA, conflictB}) + nilCount, duplicateErrCount := 0, 0 + for _, err := range conflictErrs { + switch { + case err == nil: + nilCount++ + case errors.Is(err, domain.ErrMessageRandomIDDuplicate): + duplicateErrCount++ + default: + t.Fatalf("conflicting race unexpected err = %v; results=%+v", err, conflictResults) + } + } + if nilCount != 1 || duplicateErrCount != 1 { + t.Fatalf("conflicting race errors = %v, want one success and one duplicate", conflictErrs) + } + assertChannelRandomReceiptCount(t, ctx, pool, conflictChannelID, owner.ID, conflictA.RandomID, 1) +} + +func TestChannelSendFingerprintSingleConnectionConflictLookupPostgres(t *testing.T) { + dsn := os.Getenv("TELESRV_TEST_POSTGRES_DSN") + if dsn == "" { + t.Skip("set TELESRV_TEST_POSTGRES_DSN to run postgres integration test") + } + setupPool := testPool(t) + setupCtx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(setupPool) + owner, err := users.Create(setupCtx, domain.User{AccessHash: 192, Phone: "+1781" + suffix + "21", FirstName: "OneConnectionOwner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + var channelIDs []int64 + t.Cleanup(func() { + cleanupCtx := context.Background() + if len(channelIDs) != 0 { + _, _ = setupPool.Exec(cleanupCtx, `DELETE FROM channels WHERE id = ANY($1::bigint[])`, channelIDs) + } + _, _ = setupPool.Exec(cleanupCtx, `DELETE FROM users WHERE id = $1`, owner.ID) + }) + + setupChannels := NewChannelStore(setupPool) + created, err := setupChannels.CreateChannel(setupCtx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, Title: "one connection " + suffix, Megagroup: true, Date: 1700120000, + }) + if err != nil { + t.Fatalf("create ordinary channel: %v", err) + } + channelIDs = append(channelIDs, created.Channel.ID) + ordinaryReq := domain.SendChannelMessageRequest{ + UserID: owner.ID, ChannelID: created.Channel.ID, RandomID: 792001, Message: "single pool exact", Date: 1700120001, + IdempotencyPreflighted: true, + } + broadcast, err := setupChannels.CreateChannel(setupCtx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, Title: "one connection mono " + suffix, Broadcast: true, Date: 1700120010, + }) + if err != nil { + t.Fatalf("create broadcast: %v", err) + } + channelIDs = append(channelIDs, broadcast.Channel.ID) + enabled, err := setupChannels.SetPaidMessagesPrice(setupCtx, owner.ID, broadcast.Channel.ID, 0, true) + if err != nil { + t.Fatalf("enable monoforum: %v", err) + } + monoID := enabled.Channel.LinkedMonoforumID + channelIDs = append(channelIDs, monoID) + + // Fixture creation itself has legacy allocator paths that require more than + // one connection. Constrain only the send/replay path under test. + cfg, err := pgxpool.ParseConfig(dsn) + if err != nil { + t.Fatalf("parse postgres config: %v", err) + } + cfg.MaxConns = 1 + cfg.MinConns = 0 + pool, err := pgxpool.NewWithConfig(context.Background(), cfg) + if err != nil { + t.Fatalf("open one-connection pool: %v", err) + } + t.Cleanup(pool.Close) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + t.Cleanup(cancel) + msgIDs := &singleConnectionMessageIDAllocator{current: make(map[int64]int)} + for _, id := range []int64{created.Channel.ID, monoID} { + var current int + if err := setupPool.QueryRow(setupCtx, `SELECT COALESCE(MAX(id), 0) FROM channel_messages WHERE channel_id=$1`, id).Scan(¤t); err != nil { + t.Fatalf("seed message allocator for channel %d: %v", id, err) + } + msgIDs.current[id] = current + } + oneConnectionStore := func() *ChannelStore { + return NewChannelStore(pool, WithChannelAllocators(nil, msgIDs)) + } + + var ordinaryResults [2]domain.SendChannelMessageResult + var ordinaryErrs [2]error + start := make(chan struct{}) + var wg sync.WaitGroup + for i := range ordinaryResults { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + ordinaryResults[i], ordinaryErrs[i] = oneConnectionStore().SendChannelMessage(ctx, ordinaryReq) + }(i) + } + close(start) + wg.Wait() + for i, err := range ordinaryErrs { + if err != nil { + t.Fatalf("one-connection ordinary result[%d] err = %v", i, err) + } + } + if ordinaryResults[0].Message.ID != ordinaryResults[1].Message.ID || ordinaryResults[0].Duplicate == ordinaryResults[1].Duplicate { + t.Fatalf("one-connection ordinary results = %+v / %+v, want same id and one duplicate", ordinaryResults[0], ordinaryResults[1]) + } + + monoReq := domain.SendMonoforumMessageRequest{ + MonoforumID: monoID, SenderUserID: owner.ID, + SavedPeer: domain.Peer{Type: domain.PeerTypeUser, ID: owner.ID}, + RandomID: 792002, Message: "single pool mono exact", Date: 1700120011, + IdempotencyPreflighted: true, + } + var monoResults [2]domain.SendChannelMessageResult + var monoErrs [2]error + start = make(chan struct{}) + for i := range monoResults { + wg.Add(1) + go func(i int) { + defer wg.Done() + <-start + monoResults[i], monoErrs[i] = oneConnectionStore().SendMonoforumMessage(ctx, monoReq) + }(i) + } + close(start) + wg.Wait() + for i, err := range monoErrs { + if err != nil { + t.Fatalf("one-connection monoforum result[%d] err = %v", i, err) + } + } + if monoResults[0].Message.ID != monoResults[1].Message.ID || monoResults[0].Duplicate == monoResults[1].Duplicate { + t.Fatalf("one-connection monoforum results = %+v / %+v, want same id and one duplicate", monoResults[0], monoResults[1]) + } +} + +type singleConnectionMessageIDAllocator struct { + mu sync.Mutex + current map[int64]int +} + +func (a *singleConnectionMessageIDAllocator) NextChannelMessageID(_ context.Context, channelID int64) (int, error) { + a.mu.Lock() + defer a.mu.Unlock() + a.current[channelID]++ + return a.current[channelID], nil +} + +func (a *singleConnectionMessageIDAllocator) CurrentChannelMessageID(_ context.Context, channelID int64) (int, error) { + a.mu.Lock() + defer a.mu.Unlock() + return a.current[channelID], nil +} + +func assertChannelRandomReceiptCount(t *testing.T, ctx context.Context, pool *pgxpool.Pool, channelID, senderUserID, randomID int64, want int) { + t.Helper() + var got int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM channel_messages WHERE channel_id=$1 AND sender_user_id=$2 AND random_id=$3`, channelID, senderUserID, randomID).Scan(&got); err != nil { + t.Fatalf("count channel random receipt: %v", err) + } + if got != want { + t.Fatalf("channel random receipt count = %d, want %d", got, want) + } +} + +func TestChannelSendFingerprintMigrationRoundTripPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + downSQL, err := deploy.Migrations.ReadFile("migrations/0078_channel_send_fingerprint.down.sql") + if err != nil { + t.Fatalf("read 0078 down: %v", err) + } + upSQL, err := deploy.Migrations.ReadFile("migrations/0078_channel_send_fingerprint.up.sql") + if err != nil { + t.Fatalf("read 0078 up: %v", err) + } + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin 0078 round trip: %v", err) + } + defer func() { _ = tx.Rollback(context.Background()) }() + if _, err := tx.Exec(ctx, string(downSQL)); err != nil { + t.Fatalf("0078 down: %v", err) + } + if _, err := tx.Exec(ctx, string(upSQL)); err != nil { + t.Fatalf("0078 up: %v", err) + } + var defaultExpr string + var constraintExists bool + if err := tx.QueryRow(ctx, ` +SELECT column_default, + EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'channel_messages_request_fingerprint_size') +FROM information_schema.columns +WHERE table_schema='public' AND table_name='channel_messages' AND column_name='request_fingerprint'`).Scan(&defaultExpr, &constraintExists); err != nil { + t.Fatalf("inspect 0078: %v", err) + } + if !strings.Contains(defaultExpr, `\x`) || !constraintExists { + t.Fatalf("0078 default=%q constraint=%v, want empty bytea rolling default + size constraint", defaultExpr, constraintExists) + } +} diff --git a/internal/store/postgres/channel_update_retention.go b/internal/store/postgres/channel_update_retention.go new file mode 100644 index 00000000..6d776117 --- /dev/null +++ b/internal/store/postgres/channel_update_retention.go @@ -0,0 +1,285 @@ +package postgres + +import ( + "context" + "errors" + "fmt" + "math" + "time" + + "github.com/jackc/pgx/v5" + + "telesrv/internal/domain" + "telesrv/internal/store/postgres/sqlcgen" +) + +const ( + channelUpdateRetentionCandidateBatch = 256 + // Keep one channel row/checkpoint hot-lock window short even when the maintenance pass has a + // large global budget. The outer seek loop may consume many chunks; this is a transaction cap, + // not a per-pass correctness cap. + channelUpdateRetentionTransactionBatch = 256 +) + +// PruneChannelUpdateEvents atomically removes a bounded contiguous prefix of one channel's durable +// event log. The retained floor advances only through complete event rows actually deleted; a target +// inside a pts_count interval leaves that row and the floor untouched. +func (s *ChannelStore) PruneChannelUpdateEvents(ctx context.Context, channelID int64, throughPts, limit int) (domain.ChannelUpdateRetentionResult, error) { + return s.pruneChannelUpdateEvents(ctx, channelID, throughPts, 0, limit) +} + +// DeleteExpiredChannelUpdateEvents performs a bounded global retention pass without OFFSET. The +// candidate seek uses (date,channel_id,pts), selects only the oldest retained row of each channel, +// then delegates deletion/floor advancement to the per-channel transactional primitive. +func (s *ChannelStore) DeleteExpiredChannelUpdateEvents(ctx context.Context, olderThan time.Duration, limit int) (int, error) { + if olderThan <= 0 { + return 0, nil + } + limit = normalizeChannelUpdateRetentionLimit(limit) + cutoff := int(time.Now().Add(-olderThan).Unix()) + deleted := 0 + excluded := make([]int64, 0) + var isolatedErrors []error + for deleted < limit { + candidateLimit := limit - deleted + if candidateLimit > channelUpdateRetentionCandidateBatch { + candidateLimit = channelUpdateRetentionCandidateBatch + } + channelIDs, err := s.expiredChannelUpdateCandidates(ctx, cutoff, candidateLimit, excluded) + if err != nil { + isolatedErrors = append(isolatedErrors, err) + return deleted, errors.Join(isolatedErrors...) + } + if len(channelIDs) == 0 { + break + } + for _, channelID := range channelIDs { + if deleted >= limit { + break + } + chunkLimit := limit - deleted + if chunkLimit > channelUpdateRetentionTransactionBatch { + chunkLimit = channelUpdateRetentionTransactionBatch + } + result, err := s.pruneChannelUpdateEvents(ctx, channelID, math.MaxInt32, cutoff, chunkLimit) + if err != nil { + // A durable-log gap/invalid row is an invariant violation for this channel, but it must + // not starve every healthy channel behind the oldest candidate. Isolate it for this pass, + // keep its floor unchanged (the tx rolled back), continue globally, then report all errors. + excluded = append(excluded, channelID) + isolatedErrors = append(isolatedErrors, fmt.Errorf("channel %d retention isolated: %w", channelID, err)) + continue + } + if result.Deleted == 0 { + // Another retention worker may have consumed this head after the seek. + // Exclude it for this pass so one raced channel cannot spin forever. + excluded = append(excluded, channelID) + continue + } + deleted += result.Deleted + } + } + return deleted, errors.Join(isolatedErrors...) +} + +// expiredChannelUpdateCandidates keeps each SQL seek bounded, while the caller loops through as +// many seeks as needed to consume the requested deletion budget. The 256 value is a fetch/page +// size, not a per-maintenance-pass correctness cap. +func (s *ChannelStore) expiredChannelUpdateCandidates(ctx context.Context, cutoff, limit int, excluded []int64) ([]int64, error) { + rows, err := s.db.Query(ctx, ` +SELECT e.channel_id +FROM channel_update_events e +LEFT JOIN channel_update_checkpoints cp ON cp.channel_id = e.channel_id +WHERE e.date < $1 + AND e.pts > COALESCE(cp.retained_through_pts, 0) + AND NOT (e.channel_id = ANY($3::bigint[])) + AND NOT EXISTS ( + SELECT 1 + FROM channel_update_events earlier + WHERE earlier.channel_id = e.channel_id + AND earlier.pts > COALESCE(cp.retained_through_pts, 0) + AND earlier.pts < e.pts + ) +ORDER BY e.date ASC, e.channel_id ASC, e.pts ASC +LIMIT $2`, cutoff, limit, excluded) + if err != nil { + return nil, fmt.Errorf("list expired channel update candidates: %w", err) + } + defer rows.Close() + channelIDs := make([]int64, 0, limit) + for rows.Next() { + var channelID int64 + if err := rows.Scan(&channelID); err != nil { + return nil, fmt.Errorf("scan expired channel update candidate: %w", err) + } + channelIDs = append(channelIDs, channelID) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate expired channel update candidates: %w", err) + } + return channelIDs, nil +} + +func (s *ChannelStore) pruneChannelUpdateEvents(ctx context.Context, channelID int64, throughPts, beforeDate, limit int) (domain.ChannelUpdateRetentionResult, error) { + if channelID == 0 || throughPts < 0 { + return domain.ChannelUpdateRetentionResult{}, domain.ErrChannelInvalid + } + limit = normalizeChannelUpdateRetentionLimit(limit) + if limit > channelUpdateRetentionTransactionBatch { + limit = channelUpdateRetentionTransactionBatch + } + var result domain.ChannelUpdateRetentionResult + err := withTx(ctx, s.db, "prune channel update events", func(tx pgx.Tx) error { + checkpoint, err := lockChannelUpdateCheckpoint(ctx, tx, channelID) + if err != nil { + return err + } + if throughPts > checkpoint.LatestPts { + throughPts = checkpoint.LatestPts + } + if throughPts <= checkpoint.RetainedThroughPts { + result.Checkpoint = checkpoint + return nil + } + + rows, err := tx.Query(ctx, ` +SELECT pts, pts_count, date +FROM channel_update_events +WHERE channel_id = $1 + AND pts > $2 + AND pts <= $3 +ORDER BY pts ASC +LIMIT $4 +FOR UPDATE`, channelID, checkpoint.RetainedThroughPts, throughPts, limit) + if err != nil { + return fmt.Errorf("list channel update prune prefix: %w", err) + } + defer rows.Close() + cursor := checkpoint.RetainedThroughPts + ptsToDelete := make([]int32, 0, limit) + for rows.Next() { + var pts, ptsCount, date int + if err := rows.Scan(&pts, &ptsCount, &date); err != nil { + return fmt.Errorf("scan channel update prune prefix: %w", err) + } + if beforeDate > 0 && date >= beforeDate { + break + } + if ptsCount <= 0 { + return fmt.Errorf("prune channel update events: channel %d has invalid pts_count=%d at pts=%d", channelID, ptsCount, pts) + } + if pts != cursor+ptsCount { + return fmt.Errorf( + "prune channel update events: channel %d has gap after pts %d: event pts=%d pts_count=%d", + channelID, cursor, pts, ptsCount, + ) + } + cursor = pts + ptsToDelete = append(ptsToDelete, int32(pts)) + } + if err := rows.Err(); err != nil { + return fmt.Errorf("iterate channel update prune prefix: %w", err) + } + rows.Close() + + if len(ptsToDelete) == 0 { + result.Checkpoint = checkpoint + return nil + } + tag, err := tx.Exec(ctx, ` +DELETE FROM channel_update_events +WHERE channel_id = $1 + AND pts = ANY($2::int[])`, channelID, ptsToDelete) + if err != nil { + return fmt.Errorf("delete channel update prune prefix: %w", err) + } + if got := int(tag.RowsAffected()); got != len(ptsToDelete) { + return fmt.Errorf("delete channel update prune prefix: deleted %d rows, expected %d", got, len(ptsToDelete)) + } + tag, err = tx.Exec(ctx, ` +UPDATE channel_update_checkpoints +SET retained_through_pts = $2, + latest_event_date = GREATEST(latest_event_date, $3), + latest_pts = GREATEST(latest_pts, $4), + updated_at = now() +WHERE channel_id = $1`, channelID, cursor, checkpoint.LatestEventDate, checkpoint.LatestPts) + if err != nil { + return fmt.Errorf("advance channel update retained floor: %w", err) + } + if tag.RowsAffected() != 1 { + return fmt.Errorf("advance channel update retained floor: checkpoint row disappeared for channel %d", channelID) + } + checkpoint.RetainedThroughPts = cursor + result = domain.ChannelUpdateRetentionResult{Checkpoint: checkpoint, Deleted: len(ptsToDelete)} + return nil + }) + if err != nil { + return domain.ChannelUpdateRetentionResult{}, err + } + return result, nil +} + +// lockChannelUpdateCheckpoint follows the channel writer lock order: channels row first, checkpoint +// second. Event insertion updates channels.pts before upserting the checkpoint, so retention cannot +// race a committed pts without its durable event/checkpoint. +func lockChannelUpdateCheckpoint(ctx context.Context, tx pgx.Tx, channelID int64) (domain.ChannelUpdateRetentionCheckpoint, error) { + var lockedChannelID int64 + if err := tx.QueryRow(ctx, ` +SELECT id +FROM channels +WHERE id = $1 +FOR UPDATE`, channelID).Scan(&lockedChannelID); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return domain.ChannelUpdateRetentionCheckpoint{}, domain.ErrChannelInvalid + } + return domain.ChannelUpdateRetentionCheckpoint{}, fmt.Errorf("lock channel for update retention: %w", err) + } + checkpoint := domain.ChannelUpdateRetentionCheckpoint{ChannelID: channelID} + if err := tx.QueryRow(ctx, ` +SELECT retained_through_pts, latest_event_date, latest_pts +FROM channel_update_checkpoints +WHERE channel_id = $1 +FOR UPDATE`, channelID).Scan( + &checkpoint.RetainedThroughPts, + &checkpoint.LatestEventDate, + &checkpoint.LatestPts, + ); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return domain.ChannelUpdateRetentionCheckpoint{}, fmt.Errorf( + "lock channel update checkpoint: invariant violation: channel %d has no retention checkpoint", + channelID, + ) + } + return domain.ChannelUpdateRetentionCheckpoint{}, fmt.Errorf("lock channel update checkpoint: %w", err) + } + return checkpoint, nil +} + +func normalizeChannelUpdateRetentionLimit(limit int) int { + if limit <= 0 || limit > domain.MaxChannelUpdateRetentionBatch { + return domain.MaxChannelUpdateRetentionBatch + } + return limit +} + +func getChannelUpdateCheckpoint(ctx context.Context, db sqlcgen.DBTX, channelID int64) (domain.ChannelUpdateRetentionCheckpoint, error) { + checkpoint := domain.ChannelUpdateRetentionCheckpoint{ChannelID: channelID} + err := db.QueryRow(ctx, ` +SELECT retained_through_pts, latest_event_date, latest_pts +FROM channel_update_checkpoints +WHERE channel_id = $1`, channelID).Scan( + &checkpoint.RetainedThroughPts, + &checkpoint.LatestEventDate, + &checkpoint.LatestPts, + ) + if errors.Is(err, pgx.ErrNoRows) { + return domain.ChannelUpdateRetentionCheckpoint{}, fmt.Errorf( + "get channel update checkpoint: invariant violation: channel %d has no retention checkpoint", + channelID, + ) + } + if err != nil { + return domain.ChannelUpdateRetentionCheckpoint{}, fmt.Errorf("get channel update checkpoint: %w", err) + } + return checkpoint, nil +} diff --git a/internal/store/postgres/channel_update_retention_integration_test.go b/internal/store/postgres/channel_update_retention_integration_test.go new file mode 100644 index 00000000..c15dd4bd --- /dev/null +++ b/internal/store/postgres/channel_update_retention_integration_test.go @@ -0,0 +1,232 @@ +package postgres + +import ( + "context" + "fmt" + "strings" + "testing" + "time" + + "telesrv/internal/domain" +) + +func TestChannelUpdateRetentionFloorDifferenceAndDirtyCheckpointPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + owner, err := users.Create(ctx, domain.User{ + AccessHash: 71, + Phone: "+1766" + suffix + "01", + FirstName: "RetentionOwner", + }) + if err != nil { + t.Fatalf("create owner: %v", err) + } + var channelID int64 + t.Cleanup(func() { + if channelID != 0 { + _, _ = pool.Exec(ctx, "DELETE FROM channels WHERE id = $1", channelID) + } + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", owner.ID) + }) + + channels := NewChannelStore(pool) + created, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, + Title: "Retention PG " + suffix, + Megagroup: true, + Date: 1_700_020_000, + }) + if err != nil { + t.Fatalf("create channel: %v", err) + } + channelID = created.Channel.ID + sent := make([]domain.SendChannelMessageResult, 0, 3) + for i := 1; i <= 3; i++ { + result, err := channels.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + UserID: owner.ID, ChannelID: channelID, RandomID: int64(7000 + i), Message: "retention", Date: 1_700_020_000 + i, + }) + if err != nil { + t.Fatalf("send message %d: %v", i, err) + } + sent = append(sent, result) + } + + pruned, err := channels.PruneChannelUpdateEvents(ctx, channelID, sent[1].Event.Pts, 100) + if err != nil { + t.Fatalf("prune channel updates: %v", err) + } + if pruned.Deleted != 3 || pruned.Checkpoint.RetainedThroughPts != sent[1].Event.Pts { + t.Fatalf("prune result = %+v, want deleted=3 floor=%d", pruned, sent[1].Event.Pts) + } + var floor, latestDate, latestPts, remaining int + if err := pool.QueryRow(ctx, ` +SELECT cp.retained_through_pts, cp.latest_event_date, cp.latest_pts, + (SELECT COUNT(*) FROM channel_update_events e WHERE e.channel_id = cp.channel_id)::int +FROM channel_update_checkpoints cp +WHERE cp.channel_id = $1`, channelID).Scan(&floor, &latestDate, &latestPts, &remaining); err != nil { + t.Fatalf("read retention checkpoint: %v", err) + } + if floor != sent[1].Event.Pts || latestDate != sent[2].Event.Date || latestPts != sent[2].Event.Pts || remaining != 1 { + t.Fatalf("checkpoint/db = floor:%d latest:%d/%d remaining:%d", floor, latestDate, latestPts, remaining) + } + + below, err := channels.ListChannelDifference(ctx, domain.ChannelDifferenceRequest{ + UserID: owner.ID, ChannelID: channelID, Pts: floor - 1, Limit: 100, + }) + if err != nil { + t.Fatalf("difference below retained floor: %v", err) + } + if !below.TooLong || below.Pts != sent[2].Event.Pts { + t.Fatalf("difference below floor = %+v, want too-long snapshot at pts %d", below, sent[2].Event.Pts) + } + atFloor, err := channels.ListChannelDifference(ctx, domain.ChannelDifferenceRequest{ + UserID: owner.ID, ChannelID: channelID, Pts: floor, Limit: 100, + }) + if err != nil { + t.Fatalf("difference at retained floor: %v", err) + } + if atFloor.TooLong || len(atFloor.Events) != 1 || atFloor.Events[0].Pts != sent[2].Event.Pts { + t.Fatalf("difference at floor = %+v, want normal incremental event pts %d", atFloor, sent[2].Event.Pts) + } + + allPruned, err := channels.PruneChannelUpdateEvents(ctx, channelID, sent[2].Event.Pts, 100) + if err != nil { + t.Fatalf("prune remaining channel update: %v", err) + } + if allPruned.Deleted != 1 { + t.Fatalf("remaining prune = %+v, want deleted=1", allPruned) + } + dirty, err := channels.ListDirtyActiveChannelsForUser(ctx, owner.ID, sent[2].Event.Date-1, 0, 10) + if err != nil { + t.Fatalf("list dirty channels after prune: %v", err) + } + if len(dirty) != 1 || dirty[0].ChannelID != channelID || dirty[0].Pts != sent[2].Event.Pts { + t.Fatalf("dirty channels after prune = %+v, want channel %d pts %d", dirty, channelID, sent[2].Event.Pts) + } +} + +func TestDeleteExpiredChannelUpdateEventsContinuesPastCandidatePagePostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + owner, err := users.Create(ctx, domain.User{ + AccessHash: time.Now().UnixNano(), + Phone: "+1767" + suffix + "01", + FirstName: "RetentionPageOwner", + }) + if err != nil { + t.Fatalf("create retention page owner: %v", err) + } + channelIDs := make([]int64, 0, 320) + t.Cleanup(func() { + if len(channelIDs) > 0 { + _, _ = pool.Exec(ctx, "DELETE FROM channels WHERE id = ANY($1::bigint[])", channelIDs) + } + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", owner.ID) + }) + + channels := NewChannelStore(pool) + for i := 0; i < 320; i++ { + created, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, + Title: fmt.Sprintf("Retention page %s/%03d", suffix, i), + Megagroup: true, + // Keep these rows ahead of ordinary developer/test data in the global seek. + Date: 1, + }) + if err != nil { + t.Fatalf("create retention candidate %d: %v", i, err) + } + channelIDs = append(channelIDs, created.Channel.ID) + } + + deleted, err := channels.DeleteExpiredChannelUpdateEvents(ctx, time.Second, len(channelIDs)) + if err != nil { + t.Fatalf("delete expired channel updates across pages: %v", err) + } + if deleted != len(channelIDs) { + t.Fatalf("deleted expired channel updates = %d, want %d (must continue after page 256)", deleted, len(channelIDs)) + } + var remaining, advanced int + if err := pool.QueryRow(ctx, ` +SELECT + (SELECT count(*) FROM channel_update_events WHERE channel_id = ANY($1::bigint[]))::int, + (SELECT count(*) FROM channel_update_checkpoints + WHERE channel_id = ANY($1::bigint[]) AND retained_through_pts = 1)::int`, channelIDs).Scan(&remaining, &advanced); err != nil { + t.Fatalf("read paged channel retention result: %v", err) + } + if remaining != 0 || advanced != len(channelIDs) { + t.Fatalf("paged retention remaining/advanced = %d/%d, want 0/%d", remaining, advanced, len(channelIDs)) + } +} + +func TestDeleteExpiredChannelUpdateEventsIsolatesGapAndContinuesPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + owner, err := NewUserStore(pool).Create(ctx, domain.User{ + AccessHash: time.Now().UnixNano(), + Phone: "+1768" + suffix + "01", + FirstName: "RetentionGapOwner", + }) + if err != nil { + t.Fatalf("create owner: %v", err) + } + channels := NewChannelStore(pool) + bad, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, Title: "Retention bad " + suffix, Megagroup: true, Date: 1, + }) + if err != nil { + t.Fatalf("create bad channel: %v", err) + } + healthy, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, Title: "Retention healthy " + suffix, Megagroup: true, Date: 2, + }) + if err != nil { + t.Fatalf("create healthy channel: %v", err) + } + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM channels WHERE id = ANY($1::bigint[])", []int64{bad.Channel.ID, healthy.Channel.ID}) + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", owner.ID) + }) + + // Deliberately model a persisted invariant violation: floor=0 but the first event ends at pts=2 + // with pts_count=1. The bad channel is the oldest global candidate and must be reported without + // preventing the healthy channel behind it from advancing. + if _, err := pool.Exec(ctx, ` +WITH moved_event AS ( + UPDATE channel_update_events SET pts = 2, date = 1 WHERE channel_id = $1 RETURNING channel_id +), moved_channel AS ( + UPDATE channels SET pts = 2 WHERE id = $1 RETURNING id +) +UPDATE channel_update_checkpoints +SET latest_pts = 2, latest_event_date = 1 +WHERE channel_id = $1 +`, bad.Channel.ID); err != nil { + t.Fatalf("inject channel retention gap: %v", err) + } + + deleted, err := channels.DeleteExpiredChannelUpdateEvents(ctx, time.Second, 10) + if err == nil || !strings.Contains(err.Error(), "has gap") { + t.Fatalf("gap retention err = %v, want reported invariant violation", err) + } + if deleted < 1 { + t.Fatalf("deleted across bad+healthy channels = %d, want at least the healthy channel's event", deleted) + } + var badFloor, badRows, healthyFloor, healthyRows int + if scanErr := pool.QueryRow(ctx, ` +SELECT + (SELECT retained_through_pts FROM channel_update_checkpoints WHERE channel_id = $1)::int, + (SELECT count(*) FROM channel_update_events WHERE channel_id = $1)::int, + (SELECT retained_through_pts FROM channel_update_checkpoints WHERE channel_id = $2)::int, + (SELECT count(*) FROM channel_update_events WHERE channel_id = $2)::int +`, bad.Channel.ID, healthy.Channel.ID).Scan(&badFloor, &badRows, &healthyFloor, &healthyRows); scanErr != nil { + t.Fatalf("read isolated retention state: %v", scanErr) + } + if badFloor != 0 || badRows != 1 || healthyFloor != 1 || healthyRows != 0 { + t.Fatalf("isolated state bad=%d/%d healthy=%d/%d, want 0/1 and 1/0", badFloor, badRows, healthyFloor, healthyRows) + } +} diff --git a/internal/store/postgres/channel_updates.go b/internal/store/postgres/channel_updates.go index 87ac7a52..331768fe 100644 --- a/internal/store/postgres/channel_updates.go +++ b/internal/store/postgres/channel_updates.go @@ -37,7 +37,11 @@ func (s *ChannelStore) ListChannelDifference(ctx context.Context, req domain.Cha Dialog: previewChannelDialog(req.UserID, channel, member), }, nil } - if channel.Pts-req.Pts > limit { + checkpoint, err := getChannelUpdateCheckpoint(ctx, s.db, req.ChannelID) + if err != nil { + return domain.ChannelDifference{}, err + } + if req.Pts < checkpoint.RetainedThroughPts || channel.Pts-req.Pts > limit { args := []any{req.ChannelID} where := "channel_id = $1 AND NOT deleted" if member.AvailableMinID > 0 { @@ -210,6 +214,30 @@ func (s *ChannelStore) MaxChannelPts(ctx context.Context, channelID int64) (int, return pts, err } +func (s *ChannelStore) MaxChannelPtsBatch(ctx context.Context, channelIDs []int64) (map[int64]int, error) { + out := make(map[int64]int, len(channelIDs)) + if len(channelIDs) == 0 { + return out, nil + } + rows, err := s.db.Query(ctx, `SELECT id, pts FROM channels WHERE id = ANY($1::bigint[])`, channelIDs) + if err != nil { + return nil, err + } + defer rows.Close() + for rows.Next() { + var channelID int64 + var pts int + if err := rows.Scan(&channelID, &pts); err != nil { + return nil, err + } + out[channelID] = pts + } + if err := rows.Err(); err != nil { + return nil, err + } + return out, nil +} + func transientChannelParticipantEvent(channelID, actorUserID int64, previous, participant domain.ChannelMember, date int) domain.ChannelUpdateEvent { return domain.ChannelUpdateEvent{ ChannelID: channelID, @@ -302,6 +330,18 @@ INSERT INTO channel_update_events ( ids, event.SenderUserID, userIDs, payload); err != nil { return fmt.Errorf("insert channel event: %w", err) } + // The checkpoint is updated in the same business transaction as the event row. Retention may + // later remove the row, but account-level dirty-channel recovery still has the latest date/pts. + if _, err := tx.Exec(ctx, ` +INSERT INTO channel_update_checkpoints ( + channel_id, retained_through_pts, latest_event_date, latest_pts +) VALUES ($1, 0, $2, $3) +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()`, event.ChannelID, event.Date, event.Pts); err != nil { + return fmt.Errorf("upsert channel update checkpoint: %w", err) + } return nil } diff --git a/internal/store/postgres/contiguous_pts_integration_test.go b/internal/store/postgres/contiguous_pts_integration_test.go index f8580061..fd6e747e 100644 --- a/internal/store/postgres/contiguous_pts_integration_test.go +++ b/internal/store/postgres/contiguous_pts_integration_test.go @@ -217,7 +217,7 @@ func TestDispatchOutboxLifecycleKeepsDurableEvents(t *testing.T) { if len(claimed) != 1 || claimed[0].TargetUserID != owner.ID || claimed[0].Pts != 1 || claimed[0].Attempts != 1 || claimed[0].ExcludeSessionID != 101 { t.Fatalf("claimed first = %+v, want owner pts=1 attempts=1", claimed) } - if err := outbox.MarkDelivered(ctx, owner.ID, claimed[0].ID); err != nil { + if err := outbox.MarkDelivered(ctx, claimed[0]); err != nil { t.Fatalf("MarkDelivered: %v", err) } if got := outboxRows(1); got != 0 { @@ -246,7 +246,7 @@ func TestDispatchOutboxLifecycleKeepsDurableEvents(t *testing.T) { if len(claimed) != 1 || claimed[0].TargetUserID != owner.ID || claimed[0].Pts != 2 || claimed[0].Attempts != 2 { t.Fatalf("claimed stale = %+v, want owner pts=2 attempts=2", claimed) } - if err := outbox.MarkFailed(ctx, owner.ID, claimed[0].ID, "temporary"); err != nil { + if err := outbox.MarkFailed(ctx, claimed[0], "temporary"); err != nil { t.Fatalf("MarkFailed temporary: %v", err) } var status string @@ -275,7 +275,8 @@ func TestDispatchOutboxLifecycleKeepsDurableEvents(t *testing.T) { `, owner.ID, claimed[0].ID); err != nil { t.Fatalf("prepare terminal failure: %v", err) } - if err := outbox.MarkFailed(ctx, owner.ID, claimed[0].ID, "permanent"); err != nil { + claimed[0].Attempts = 5 + if err := outbox.MarkFailed(ctx, claimed[0], "permanent"); err != nil { t.Fatalf("MarkFailed permanent: %v", err) } if err := tx.QueryRow(ctx, ` diff --git a/internal/store/postgres/dispatch_outbox.go b/internal/store/postgres/dispatch_outbox.go index 7f54baa1..aeb9bad5 100644 --- a/internal/store/postgres/dispatch_outbox.go +++ b/internal/store/postgres/dispatch_outbox.go @@ -12,7 +12,11 @@ import ( // defaultDispatchLease 是 'dispatching' 行被判定租约过期、可被重新 claim 的默认时长。 // 与 docs/message-module.md 的 outbox 背压参数对应；生产由 config 注入覆盖。 -const defaultDispatchLease = 30 * time.Second +const ( + defaultDispatchLease = 30 * time.Second + defaultDispatchPoisonCleanupBatch = 256 + maxDispatchPoisonCleanupBatch = 1000 +) // DispatchOutboxStore 用 PostgreSQL 实现 transactional outbox。 type DispatchOutboxStore struct { @@ -63,6 +67,48 @@ func (s *DispatchOutboxStore) ClaimPending(ctx context.Context, limit int) ([]st if err != nil { return nil, fmt.Errorf("claim dispatch outbox: %w", err) } + return dispatchItemsFromClaimRows(rows), nil +} + +// ClaimPendingShards 只领取固定 logical shard 集合中的用户 head 事件。 +// shardCount 是稳定哈希空间，shardIDs 是当前 worker 独占的子集；worker 数变化只改变 +// shard→worker 的运行时归属，不改变 user→shard，从而避免同一用户被并行领取。 +func (s *DispatchOutboxStore) ClaimPendingShards(ctx context.Context, shardCount int, shardIDs []int, limit int) ([]store.DispatchOutboxItem, error) { + if shardCount <= 0 || len(shardIDs) == 0 { + return nil, nil + } + if shardCount != store.DispatchOutboxLogicalShards { + return nil, fmt.Errorf("claim dispatch outbox shards: shard count %d, want stable %d", shardCount, store.DispatchOutboxLogicalShards) + } + if limit <= 0 { + limit = 100 + } + if limit > 1000 { + limit = 1000 + } + ids := make([]int16, 0, len(shardIDs)) + seen := make(map[int]struct{}, len(shardIDs)) + for _, id := range shardIDs { + if id < 0 || id >= shardCount { + continue + } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + ids = append(ids, int16(id)) + } + if len(ids) == 0 { + return nil, nil + } + rows, err := s.q.ClaimDispatchOutboxShards(ctx, sqlcgen.ClaimDispatchOutboxShardsParams{ + LeaseSeconds: s.leaseSeconds, + LimitCount: int32(limit), + ShardIds: ids, + }) + if err != nil { + return nil, fmt.Errorf("claim dispatch outbox shards: %w", err) + } out := make([]store.DispatchOutboxItem, 0, len(rows)) for _, row := range rows { out = append(out, store.DispatchOutboxItem{ @@ -78,6 +124,22 @@ func (s *DispatchOutboxStore) ClaimPending(ctx context.Context, limit int) ([]st return out, nil } +func dispatchItemsFromClaimRows(rows []sqlcgen.ClaimDispatchOutboxRow) []store.DispatchOutboxItem { + out := make([]store.DispatchOutboxItem, 0, len(rows)) + for _, row := range rows { + out = append(out, store.DispatchOutboxItem{ + ID: row.ID, + TargetUserID: row.TargetUserID, + Pts: int(row.Pts), + EventType: domain.UpdateEventType(row.EventType), + ExcludeAuthKeyID: authKeyIDFromInt64(row.ExcludeAuthKeyID), + ExcludeSessionID: row.ExcludeSessionID, + Attempts: int(row.Attempts), + }) + } + return out +} + // MarkDeliveredBatch 一次性删除一批已投递的 outbox 行（方案 A：投递成功即删），取代逐条 MarkDelivered。 func (s *DispatchOutboxStore) MarkDeliveredBatch(ctx context.Context, items []store.DispatchOutboxItem) error { if len(items) == 0 { @@ -85,49 +147,66 @@ func (s *DispatchOutboxStore) MarkDeliveredBatch(ctx context.Context, items []st } targetUserIDs := make([]int64, len(items)) ids := make([]int64, len(items)) + expectedAttempts := make([]int32, len(items)) for i, it := range items { targetUserIDs[i] = it.TargetUserID ids[i] = it.ID + expectedAttempts[i] = int32(it.Attempts) } - if err := s.q.MarkDispatchDeliveredBatch(ctx, sqlcgen.MarkDispatchDeliveredBatchParams{ - TargetUserIds: targetUserIDs, - Ids: ids, - }); err != nil { + rows, err := s.q.MarkDispatchDeliveredBatch(ctx, sqlcgen.MarkDispatchDeliveredBatchParams{ + TargetUserIds: targetUserIDs, + Ids: ids, + ExpectedAttempts: expectedAttempts, + }) + if err != nil { return fmt.Errorf("mark dispatch delivered batch: %w", err) } - return nil -} - -func (s *DispatchOutboxStore) MarkDelivered(ctx context.Context, targetUserID, id int64) error { - if err := s.q.MarkDispatchDelivered(ctx, sqlcgen.MarkDispatchDeliveredParams{ - TargetUserID: targetUserID, - ID: id, - }); err != nil { - return fmt.Errorf("mark dispatch delivered: %w", err) + if rows != int64(len(items)) { + return fmt.Errorf("mark dispatch delivered batch: %w: updated %d of %d", store.ErrDispatchLeaseLost, rows, len(items)) } return nil } -func (s *DispatchOutboxStore) MarkFailed(ctx context.Context, targetUserID, id int64, lastError string) error { - if err := s.q.MarkDispatchFailed(ctx, sqlcgen.MarkDispatchFailedParams{ - TargetUserID: targetUserID, - ID: id, - LastError: lastError, - }); err != nil { +func (s *DispatchOutboxStore) MarkDelivered(ctx context.Context, item store.DispatchOutboxItem) error { + rows, err := s.q.MarkDispatchDelivered(ctx, sqlcgen.MarkDispatchDeliveredParams{ + TargetUserID: item.TargetUserID, + ID: item.ID, + ExpectedAttempts: int32(item.Attempts), + }) + if err != nil { + return fmt.Errorf("mark dispatch delivered: %w", err) + } + if rows != 1 { + return fmt.Errorf("mark dispatch delivered: %w", store.ErrDispatchLeaseLost) + } + return nil +} + +func (s *DispatchOutboxStore) MarkFailed(ctx context.Context, item store.DispatchOutboxItem, lastError string) error { + rows, err := s.q.MarkDispatchFailed(ctx, sqlcgen.MarkDispatchFailedParams{ + TargetUserID: item.TargetUserID, + ID: item.ID, + LastError: lastError, + ExpectedAttempts: int32(item.Attempts), + }) + if err != nil { return fmt.Errorf("mark dispatch failed: %w", err) } + if rows != 1 { + return fmt.Errorf("mark dispatch failed: %w", store.ErrDispatchLeaseLost) + } return nil } func (s *DispatchOutboxStore) DeleteFailed(ctx context.Context, olderThan time.Duration, limit int) (int, error) { if olderThan <= 0 { - olderThan = 24 * time.Hour + olderThan = time.Minute } if limit <= 0 { - limit = 10000 + limit = defaultDispatchPoisonCleanupBatch } - if limit > 100000 { - limit = 100000 + if limit > maxDispatchPoisonCleanupBatch { + limit = maxDispatchPoisonCleanupBatch } deleted, err := s.q.DeleteFailedDispatchOutbox(ctx, sqlcgen.DeleteFailedDispatchOutboxParams{ OlderThanSeconds: int32(olderThan / time.Second), diff --git a/internal/store/postgres/dispatch_outbox_sharding_integration_test.go b/internal/store/postgres/dispatch_outbox_sharding_integration_test.go new file mode 100644 index 00000000..81c0cf0d --- /dev/null +++ b/internal/store/postgres/dispatch_outbox_sharding_integration_test.go @@ -0,0 +1,352 @@ +package postgres + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "telesrv/internal/domain" + storepkg "telesrv/internal/store" +) + +func TestDispatchOutboxUserHeadBlocksHigherPts(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + owner := createTestUser(t, ctx, NewUserStore(pool), "+1884"+suffix+"01", "OutboxHead", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", owner.ID) + }) + + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin: %v", err) + } + defer func() { _ = tx.Rollback(ctx) }() + // ClaimPending is intentionally global. Isolate this transaction from durable tasks left by + // earlier integration tests; the rollback restores those rows after this case completes. + if _, err := tx.Exec(ctx, `DELETE FROM dispatch_outbox`); err != nil { + t.Fatalf("isolate dispatch outbox: %v", err) + } + events := NewUpdateEventStore(tx) + outbox := NewDispatchOutboxStore(tx, WithLeaseTimeout(time.Hour)) + appendEvent := func() int { + t.Helper() + event, err := events.AppendAllocatedWithDispatch(ctx, owner.ID, domain.UpdateEvent{ + Type: domain.UpdateEventDialogPinned, + PtsCount: 1, + Date: 1700002000, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: owner.ID}, + Bool: true, + }, [8]byte{}, 0) + if err != nil { + t.Fatalf("append event: %v", err) + } + return event.Pts + } + shard := int(owner.ID % int64(storepkg.DispatchOutboxLogicalShards)) + + pts1, pts2 := appendEvent(), appendEvent() + assertDispatchHead := func(wantPts int) { + t.Helper() + var gotPts int + err := tx.QueryRow(ctx, ` +SELECT head_pts +FROM dispatch_outbox_user_heads +WHERE target_user_id = $1 +`, owner.ID).Scan(&gotPts) + if err != nil { + t.Fatalf("load durable dispatch head: %v", err) + } + if gotPts != wantPts { + t.Fatalf("durable dispatch head pts = %d, want %d", gotPts, wantPts) + } + } + assertDispatchHead(pts1) + wrongShard := (shard + 1) % storepkg.DispatchOutboxLogicalShards + if wrong, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards, []int{wrongShard}, 100); err != nil || len(wrong) != 0 { + t.Fatalf("wrong-shard claim = %+v err=%v, want empty", wrong, err) + } + claimed, err := outbox.ClaimPending(ctx, 100) + if err != nil { + t.Fatalf("claim head: %v", err) + } + if len(claimed) != 1 || claimed[0].TargetUserID != owner.ID || claimed[0].Pts != pts1 { + t.Fatalf("first claim = %+v, want only pts %d", claimed, pts1) + } + if blocked, err := outbox.ClaimPending(ctx, 100); err != nil || len(blocked) != 0 { + t.Fatalf("claim behind live dispatching head = %+v err=%v, want empty (pts %d blocked)", blocked, err, pts2) + } + if blocked, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards, []int{shard}, 100); err != nil || len(blocked) != 0 { + t.Fatalf("shard claim behind live dispatching head = %+v err=%v, want empty", blocked, err) + } + if _, err := tx.Exec(ctx, `UPDATE dispatch_outbox SET updated_at = now() - interval '2 hours' WHERE target_user_id = $1 AND id = $2`, owner.ID, claimed[0].ID); err != nil { + t.Fatalf("age dispatch lease: %v", err) + } + reclaimed, err := outbox.ClaimPending(ctx, 100) + if err != nil { + t.Fatalf("reclaim stale head: %v", err) + } + if len(reclaimed) != 1 || reclaimed[0].Pts != pts1 || reclaimed[0].Attempts != 2 { + t.Fatalf("stale reclaim = %+v, want pts %d attempts 2", reclaimed, pts1) + } + if err := outbox.MarkDelivered(ctx, claimed[0]); !errors.Is(err, storepkg.ErrDispatchLeaseLost) { + t.Fatalf("old lease delivered err = %v, want ErrDispatchLeaseLost", err) + } + if err := outbox.MarkFailed(ctx, claimed[0], "stale worker"); !errors.Is(err, storepkg.ErrDispatchLeaseLost) { + t.Fatalf("old lease failed err = %v, want ErrDispatchLeaseLost", err) + } + var fencedStatus string + var fencedAttempts int + if err := tx.QueryRow(ctx, `SELECT status, attempts FROM dispatch_outbox WHERE target_user_id = $1 AND id = $2`, owner.ID, reclaimed[0].ID).Scan(&fencedStatus, &fencedAttempts); err != nil { + t.Fatalf("load fenced head: %v", err) + } + if fencedStatus != "dispatching" || fencedAttempts != 2 { + t.Fatalf("fenced head = status %s attempts %d, want dispatching/2", fencedStatus, fencedAttempts) + } + if err := outbox.MarkDelivered(ctx, reclaimed[0]); err != nil { + t.Fatalf("deliver head: %v", err) + } + assertDispatchHead(pts2) + next, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards, []int{shard}, 100) + if err != nil { + t.Fatalf("claim next after head delivered: %v", err) + } + if len(next) != 1 || next[0].Pts != pts2 { + t.Fatalf("next claim = %+v, want pts %d", next, pts2) + } + if err := outbox.MarkDelivered(ctx, next[0]); err != nil { + t.Fatalf("deliver second: %v", err) + } + var remainingHeads int + if err := tx.QueryRow(ctx, `SELECT count(*) FROM dispatch_outbox_user_heads WHERE target_user_id = $1`, owner.ID).Scan(&remainingHeads); err != nil { + t.Fatalf("count durable dispatch heads: %v", err) + } + if remainingHeads != 0 { + t.Fatalf("durable dispatch heads after lane drain = %d, want 0", remainingHeads) + } + + pts3, pts4 := appendEvent(), appendEvent() + head, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards, []int{shard}, 100) + if err != nil { + t.Fatalf("claim terminal-failure head: %v", err) + } + if len(head) != 1 || head[0].Pts != pts3 { + t.Fatalf("terminal head = %+v, want pts %d", head, pts3) + } + if _, err := tx.Exec(ctx, `UPDATE dispatch_outbox SET status = 'dispatching', attempts = 5 WHERE target_user_id = $1 AND id = $2`, owner.ID, head[0].ID); err != nil { + t.Fatalf("prepare terminal failure: %v", err) + } + head[0].Attempts = 5 + if err := outbox.MarkFailed(ctx, head[0], "permanent"); err != nil { + t.Fatalf("mark terminal failed: %v", err) + } + if got, err := outbox.ClaimPending(ctx, 100); err != nil || len(got) != 0 { + t.Fatalf("global claim behind failed head = %+v err=%v, want empty (pts %d blocked)", got, err, pts4) + } + if got, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards, []int{shard}, 100); err != nil || len(got) != 0 { + t.Fatalf("shard claim behind failed head = %+v err=%v, want empty", got, err) + } + if _, err := tx.Exec(ctx, `UPDATE dispatch_outbox SET updated_at = now() - interval '2 minutes' WHERE target_user_id = $1 AND id = $2`, owner.ID, head[0].ID); err != nil { + t.Fatalf("age poison head: %v", err) + } + if deleted, err := outbox.DeleteFailed(ctx, time.Minute, 1); err != nil || deleted != 1 { + t.Fatalf("delete quarantined failed head = %d err=%v, want 1", deleted, err) + } + var durablePoisonEvent int + if err := tx.QueryRow(ctx, `SELECT count(*) FROM user_update_events WHERE user_id = $1 AND pts = $2`, owner.ID, pts3).Scan(&durablePoisonEvent); err != nil || durablePoisonEvent != 1 { + t.Fatalf("durable poison event count = %d err=%v, want 1 for difference recovery", durablePoisonEvent, err) + } + assertDispatchHead(pts4) + unblocked, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards, []int{shard}, 100) + if err != nil { + t.Fatalf("claim after failed cleanup: %v", err) + } + if len(unblocked) != 1 || unblocked[0].Pts != pts4 { + t.Fatalf("claim after failed cleanup = %+v, want pts %d", unblocked, pts4) + } + if err := outbox.MarkDelivered(ctx, unblocked[0]); err != nil { + t.Fatalf("deliver unblocked: %v", err) + } + + pts5 := appendEvent() + tag, err := tx.Exec(ctx, ` +INSERT INTO dispatch_outbox (target_user_id, pts, event_type) +VALUES ($1, $2, $3) +ON CONFLICT DO NOTHING +`, owner.ID, pts5, string(domain.UpdateEventDialogPinned)) + if err != nil { + t.Fatalf("duplicate enqueue: %v", err) + } + if tag.RowsAffected() != 0 { + t.Fatalf("duplicate enqueue rows = %d, want 0 from (user,pts) unique key", tag.RowsAffected()) + } + var taskCount int + if err := tx.QueryRow(ctx, `SELECT count(*) FROM dispatch_outbox WHERE target_user_id = $1 AND pts = $2`, owner.ID, pts5).Scan(&taskCount); err != nil || taskCount != 1 { + t.Fatalf("duplicate task count = %d err=%v, want 1", taskCount, err) + } + if _, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards-1, []int{shard}, 1); err == nil { + t.Fatal("unstable shard count accepted") + } +} + +func TestDispatchOutboxShardClaimersAreMutuallyExclusive(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + // This case claims through the real pool from two concurrent transactions. Clear stale tasks + // left by unrelated cases so the assertion measures this one user lane, not suite order. + if _, err := pool.Exec(ctx, `DELETE FROM dispatch_outbox`); err != nil { + t.Fatalf("isolate dispatch outbox: %v", err) + } + suffix := randomSuffix(t) + owner := createTestUser(t, ctx, NewUserStore(pool), "+1885"+suffix+"01", "OutboxLane", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM dispatch_outbox WHERE target_user_id = $1", owner.ID) + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", owner.ID) + }) + if _, err := NewUpdateEventStore(pool).AppendAllocatedWithDispatch(ctx, owner.ID, domain.UpdateEvent{ + Type: domain.UpdateEventDialogPinned, + PtsCount: 1, + Date: 1700002100, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: owner.ID}, + Bool: true, + }, [8]byte{}, 0); err != nil { + t.Fatalf("append event: %v", err) + } + + outbox := NewDispatchOutboxStore(pool, WithLeaseTimeout(time.Hour)) + shard := int(owner.ID % int64(storepkg.DispatchOutboxLogicalShards)) + start := make(chan struct{}) + results := make(chan []storepkg.DispatchOutboxItem, 2) + errs := make(chan error, 2) + var wg sync.WaitGroup + for range 2 { + wg.Add(1) + go func() { + defer wg.Done() + <-start + items, err := outbox.ClaimPendingShards(ctx, storepkg.DispatchOutboxLogicalShards, []int{shard}, 1) + if err != nil { + errs <- err + return + } + results <- items + }() + } + close(start) + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("concurrent shard claim: %v", err) + } + claimed := 0 + for items := range results { + claimed += len(items) + } + if claimed != 1 { + t.Fatalf("concurrent claimed rows = %d, want exactly one user head", claimed) + } +} + +func TestDispatchOutboxLeaseExpiryAndBatchCompletionShareLockOrderPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + if _, err := pool.Exec(ctx, `DELETE FROM dispatch_outbox`); err != nil { + t.Fatalf("isolate dispatch outbox: %v", err) + } + suffix := randomSuffix(t) + users := NewUserStore(pool) + first := createTestUser(t, ctx, users, "+1886"+suffix+"01", "OutboxLockA", "") + second := createTestUser(t, ctx, users, "+1886"+suffix+"02", "OutboxLockB", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, `DELETE FROM dispatch_outbox WHERE target_user_id = ANY($1::bigint[])`, []int64{first.ID, second.ID}) + _, _ = pool.Exec(ctx, `DELETE FROM users WHERE id = ANY($1::bigint[])`, []int64{first.ID, second.ID}) + }) + events := NewUpdateEventStore(pool) + outbox := NewDispatchOutboxStore(pool, WithLeaseTimeout(time.Second)) + + for round := 0; round < 20; round++ { + for _, userID := range []int64{first.ID, second.ID} { + if _, err := events.AppendAllocatedWithDispatch(ctx, userID, domain.UpdateEvent{ + Type: domain.UpdateEventDialogPinned, PtsCount: 1, Date: 1_700_030_000 + round, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: userID}, Bool: true, + }, [8]byte{}, 0); err != nil { + t.Fatalf("round %d append user %d: %v", round, userID, err) + } + } + claimed, err := outbox.ClaimPending(ctx, 2) + if err != nil || len(claimed) != 2 { + t.Fatalf("round %d initial claim = %+v err=%v, want 2", round, claimed, err) + } + if _, err := pool.Exec(ctx, ` +UPDATE dispatch_outbox +SET updated_at = now() - interval '2 seconds' +WHERE (target_user_id, id) IN (($1, $2), ($3, $4)) +`, claimed[0].TargetUserID, claimed[0].ID, claimed[1].TargetUserID, claimed[1].ID); err != nil { + t.Fatalf("round %d age leases: %v", round, err) + } + reversed := []storepkg.DispatchOutboxItem{claimed[1], claimed[0]} + start := make(chan struct{}) + errs := make(chan error, 2) + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + <-start + _, claimErr := outbox.ClaimPending(ctx, 2) + errs <- claimErr + }() + go func() { + defer wg.Done() + <-start + markErr := outbox.MarkDeliveredBatch(ctx, reversed) + if errors.Is(markErr, storepkg.ErrDispatchLeaseLost) { + markErr = nil + } + errs <- markErr + }() + close(start) + wg.Wait() + close(errs) + for raceErr := range errs { + if raceErr != nil { + t.Fatalf("round %d lease/completion race: %v", round, raceErr) + } + } + if _, err := pool.Exec(ctx, `DELETE FROM dispatch_outbox WHERE target_user_id = ANY($1::bigint[])`, []int64{first.ID, second.ID}); err != nil { + t.Fatalf("round %d drain raced rows: %v", round, err) + } + } +} + +func TestDispatchOutboxDurableHeadRejectsStaleRowReference(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + owner := createTestUser(t, ctx, NewUserStore(pool), "+1886"+suffix+"01", "OutboxHeadFK", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM dispatch_outbox WHERE target_user_id = $1", owner.ID) + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = $1", owner.ID) + }) + if _, err := NewUpdateEventStore(pool).AppendAllocatedWithDispatch(ctx, owner.ID, domain.UpdateEvent{ + Type: domain.UpdateEventDialogPinned, PtsCount: 1, Date: 1700002200, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: owner.ID}, Bool: true, + }, [8]byte{}, 0); err != nil { + t.Fatalf("append event: %v", err) + } + + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin: %v", err) + } + defer func() { _ = tx.Rollback(ctx) }() + if _, err := tx.Exec(ctx, `UPDATE dispatch_outbox_user_heads SET head_id = head_id + 1000000 WHERE target_user_id = $1`, owner.ID); err != nil { + t.Fatalf("stage stale head: %v", err) + } + if _, err := tx.Exec(ctx, `SET CONSTRAINTS dispatch_outbox_user_heads_outbox_fkey IMMEDIATE`); err == nil { + t.Fatal("stale durable head reference unexpectedly satisfied deferred FK") + } +} diff --git a/internal/store/postgres/groupcall_integration_test.go b/internal/store/postgres/groupcall_integration_test.go index 113275da..799499df 100644 --- a/internal/store/postgres/groupcall_integration_test.go +++ b/internal/store/postgres/groupcall_integration_test.go @@ -15,13 +15,25 @@ import ( func TestGroupCallStoreContractPostgres(t *testing.T) { pool := testPool(t) ctx := context.Background() - var nextChannel int64 = 910_000_000 + // A run-unique namespace keeps filtered subtest runs independent from stale + // rows left by an interrupted older run. Per-subtest cleanup below still makes + // successful runs leave no state behind. + var nextChannel = int64(1_000_000_000 + time.Now().UnixNano()%100_000_000) storetest.RunGroupCallStoreContract(t, func(t *testing.T) (store.GroupCallStore, int64) { nextChannel++ channelID := nextChannel - t.Cleanup(func() { - _, _ = pool.Exec(ctx, "DELETE FROM group_calls WHERE channel_id = $1", channelID) - }) + // Conference contract rows have channel_id=0 and derive their call IDs from + // this synthetic channel namespace. Clean both shapes before and after each + // subtest so a previous failed run cannot feed discarded calls into the next + // run (and so conference invite/chain rows cascade away as well). + cleanup := func() { + _, _ = pool.Exec(ctx, ` +DELETE FROM group_calls +WHERE channel_id = $1 + OR (call_id >= $1 * 100 AND call_id < $1 * 100 + 100)`, channelID) + } + cleanup() + t.Cleanup(cleanup) return NewGroupCallStore(pool), channelID }) } diff --git a/internal/store/postgres/login_code_delivery.go b/internal/store/postgres/login_code_delivery.go new file mode 100644 index 00000000..827373a1 --- /dev/null +++ b/internal/store/postgres/login_code_delivery.go @@ -0,0 +1,342 @@ +package postgres + +import ( + "context" + "encoding/binary" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" + + "telesrv/internal/domain" + "telesrv/internal/store" + "telesrv/internal/store/postgres/sqlcgen" +) + +// The two-int advisory-lock namespace is disjoint from the one-bigint user +// locks used by lockUsersForUpdate. Only 32 digest bits are needed here: +// collisions merely serialize unrelated deliveries and cannot merge receipts. +const loginCodeDeliveryAdvisoryNamespace int32 = 0x4c434f44 // "LCOD" + +const ( + loginCodeDeliveryRecoveryTimeout = 2 * time.Second + loginCodeDeliveryRecoveryPoll = 20 * time.Millisecond +) + +type loginCodeDeliveryReceiptQuerier interface { + QueryRow(ctx context.Context, sql string, args ...any) pgx.Row +} + +type loginCodeDeliveryReceipt struct { + userID int64 + codeFingerprint []byte + privateMessageID int64 + messageBoxID int + pts int + messageDate int +} + +// DeliverLoginCodeMessage commits the account-visible 777000 message, dialog +// projection, user pts event, dispatch outbox row and compact idempotency +// receipt in one transaction. The raw phone_code_hash is never persisted. +func (s *MessageStore) DeliverLoginCodeMessage(ctx context.Context, req domain.LoginCodeDeliveryRequest) (domain.LoginCodeDeliveryResult, error) { + deliveryKey, err := store.LoginCodeDeliveryKey(req.PhoneCodeHash) + if err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + codeFingerprint, err := store.LoginCodeFingerprint(req.PhoneCodeHash, req.Code) + if err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + if req.Date == 0 { + req.Date = int(time.Now().Unix()) + } + if req.ExpiresAt <= int64(req.Date) { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("login code receipt expiry: %w: date=%d expires_at=%d", domain.ErrLoginCodeDeliveryInvalid, req.Date, req.ExpiresAt) + } + base, err := domain.OfficialLoginCodeMessage(req.UserID, req.Code, req.Date) + if err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + entitiesJSON, err := encodeMessageEntities(base.Entities) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("encode login code entities: %w", err) + } + + beginner, ok := s.db.(txBeginner) + if !ok { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("deliver login code: database does not support transactions") + } + tx, err := beginner.Begin(ctx) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("begin login code delivery: %w", err) + } + committed := false + defer func() { + if !committed { + rollbackCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), loginCodeDeliveryRecoveryTimeout) + defer cancel() + _ = tx.Rollback(rollbackCtx) + } + }() + + // Serialize the global idempotency key before any per-user row/advisory + // lock. This makes same-key concurrent calls deterministic even if a caller + // accidentally supplies a different user ID. + lockKey := int32(binary.BigEndian.Uint32(deliveryKey[:4])) + if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1::integer, $2::integer)`, loginCodeDeliveryAdvisoryNamespace, lockKey); err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("lock login code delivery: %w", err) + } + + receipt, found, err := getLoginCodeDeliveryReceipt(ctx, tx, deliveryKey) + if err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + if found { + if receipt.userID != req.UserID || !store.SameLoginCodeFingerprint(receipt.codeFingerprint, codeFingerprint) { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("deliver login code replay: %w", domain.ErrLoginCodeDeliveryConflict) + } + msg, err := store.RestoreLoginCodeDeliveryMessage( + receipt.userID, + req.Code, + receipt.messageDate, + receipt.privateMessageID, + receipt.messageBoxID, + receipt.pts, + ) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("restore login code replay: %w", err) + } + return domain.LoginCodeDeliveryResult{Message: msg, Created: false}, nil + } + + // All user-scoped message/update writers share this lock and acquire it + // before watermark/dialog rows, keeping box IDs and pts contiguous. + if err := lockUsersForUpdate(ctx, tx, req.UserID); err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("lock login code recipient: %w", err) + } + if err := ensureOfficialSystemUserWithDB(ctx, tx, base); err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + qtx := sqlcgen.New(tx) + + pm, err := qtx.CreatePrivateMessage(ctx, sqlcgen.CreatePrivateMessageParams{ + SenderUserID: domain.OfficialSystemUserID, + RecipientUserID: req.UserID, + RandomID: 0, + MessageDate: int32(base.Date), + Body: base.Body, + RequestFingerprint: []byte{}, + RecipientDelivered: true, + EntitiesJson: entitiesJSON, + QuoteEntitiesJson: []byte("[]"), + MediaJson: []byte("{}"), + ReplyMarkupJson: []byte("{}"), + RichMessageJson: []byte("{}"), + }) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("create login code private message: %w", err) + } + + boxID, err := s.nextLoginCodeBoxID(ctx, qtx, req.UserID) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("allocate login code box id: %w", err) + } + if boxID <= 0 || boxID > domain.MaxMessageBoxID { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("allocate login code box id: %w: %d", domain.ErrLoginCodeDeliveryInvalid, boxID) + } + pts, err := s.reservePts(ctx, tx, req.UserID) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("allocate login code pts: %w", err) + } + + boxRow, err := qtx.CreateMessageBox(ctx, sqlcgen.CreateMessageBoxParams{ + OwnerUserID: req.UserID, + BoxID: int32(boxID), + PrivateMessageID: pm.ID, + MessageSenderID: domain.OfficialSystemUserID, + PeerType: string(domain.PeerTypeUser), + PeerID: domain.OfficialSystemUserID, + FromUserID: domain.OfficialSystemUserID, + MessageDate: int32(base.Date), + Outgoing: false, + Body: base.Body, + EntitiesJson: entitiesJSON, + QuoteEntitiesJson: []byte("[]"), + Pts: int32(pts), + MediaJson: []byte("{}"), + ReplyMarkupJson: []byte("{}"), + RichMessageJson: []byte("{}"), + }) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("create login code recipient box: %w", err) + } + msg := messageFromBoxRow(boxRow) + + if err := qtx.UpsertInboxDialog(ctx, sqlcgen.UpsertInboxDialogParams{ + UserID: req.UserID, + PeerType: string(domain.PeerTypeUser), + PeerID: domain.OfficialSystemUserID, + TopMessageID: int32(msg.ID), + TopMessageDate: int32(msg.Date), + }); err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("upsert login code dialog: %w", err) + } + if err := appendNewMessageEvent(ctx, qtx, msg); err != nil { + return domain.LoginCodeDeliveryResult{}, err + } + if err := qtx.EnqueueDispatch(ctx, sqlcgen.EnqueueDispatchParams{ + TargetUserID: req.UserID, + Pts: int32(msg.Pts), + EventType: string(domain.UpdateEventNewMessage), + ExcludeAuthKeyID: 0, + ExcludeSessionID: 0, + }); err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("enqueue login code dispatch: %w", err) + } + + tag, err := tx.Exec(ctx, ` +UPDATE private_messages +SET recipient_box_id = $3, + recipient_pts = $4 +WHERE sender_user_id = $1 + AND id = $2 + AND recipient_delivered + AND recipient_box_id = 0 + AND recipient_pts = 0`, domain.OfficialSystemUserID, pm.ID, msg.ID, msg.Pts) + if err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("save login code private receipt: %w", err) + } + if tag.RowsAffected() != 1 { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("save login code private receipt: message %d lost its allocation boundary", pm.ID) + } + + if _, err := tx.Exec(ctx, ` +INSERT INTO login_code_message_deliveries ( + delivery_key, + code_fingerprint, + user_id, + private_message_id, + message_box_id, + pts, + message_date, + expires_at +) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)`, + deliveryKey[:], codeFingerprint[:], req.UserID, msg.UID, msg.ID, msg.Pts, msg.Date, time.Unix(req.ExpiresAt, 0).UTC(), + ); err != nil { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("save login code delivery receipt: %w", err) + } + + if err := tx.Commit(ctx); err != nil { + result, recoverErr := s.recoverLoginCodeDeliveryAfterCommitError(ctx, req, deliveryKey, codeFingerprint) + if recoverErr != nil { + return domain.LoginCodeDeliveryResult{}, errors.Join( + fmt.Errorf("commit login code delivery: %w", err), + recoverErr, + ) + } + committed = true + return result, nil + } + committed = true + return domain.LoginCodeDeliveryResult{Message: msg, Created: true}, nil +} + +func (s *MessageStore) recoverLoginCodeDeliveryAfterCommitError( + ctx context.Context, + req domain.LoginCodeDeliveryRequest, + deliveryKey, codeFingerprint [32]byte, +) (domain.LoginCodeDeliveryResult, error) { + probeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), loginCodeDeliveryRecoveryTimeout) + defer cancel() + ticker := time.NewTicker(loginCodeDeliveryRecoveryPoll) + defer ticker.Stop() + for { + receipt, found, err := getLoginCodeDeliveryReceipt(probeCtx, s.db, deliveryKey) + if err != nil { + return domain.LoginCodeDeliveryResult{}, errors.Join( + domain.ErrLoginCodeDeliveryCommitAmbiguous, + fmt.Errorf("probe login code delivery receipt after commit error: %w", err), + ) + } + if found { + if receipt.userID != req.UserID || !store.SameLoginCodeFingerprint(receipt.codeFingerprint, codeFingerprint) { + return domain.LoginCodeDeliveryResult{}, fmt.Errorf("probe login code delivery receipt after commit error: %w", domain.ErrLoginCodeDeliveryConflict) + } + msg, err := store.RestoreLoginCodeDeliveryMessage( + receipt.userID, + req.Code, + receipt.messageDate, + receipt.privateMessageID, + receipt.messageBoxID, + receipt.pts, + ) + if err != nil { + return domain.LoginCodeDeliveryResult{}, errors.Join( + domain.ErrLoginCodeDeliveryCommitAmbiguous, + fmt.Errorf("restore probed login code delivery: %w", err), + ) + } + // The receipt proves durable success but cannot prove whether this + // caller or an equivalent replay won the commit race. + return domain.LoginCodeDeliveryResult{Message: msg, Created: false}, nil + } + select { + case <-probeCtx.Done(): + return domain.LoginCodeDeliveryResult{}, errors.Join( + domain.ErrLoginCodeDeliveryCommitAmbiguous, + fmt.Errorf("probe login code delivery receipt after commit error: %w", probeCtx.Err()), + ) + case <-ticker.C: + } + } +} + +func getLoginCodeDeliveryReceipt(ctx context.Context, q loginCodeDeliveryReceiptQuerier, deliveryKey [32]byte) (loginCodeDeliveryReceipt, bool, error) { + var receipt loginCodeDeliveryReceipt + var boxID, pts, messageDate int32 + err := q.QueryRow(ctx, ` +SELECT user_id, + code_fingerprint, + private_message_id, + message_box_id, + pts, + message_date +FROM login_code_message_deliveries +WHERE delivery_key = $1`, deliveryKey[:]).Scan( + &receipt.userID, + &receipt.codeFingerprint, + &receipt.privateMessageID, + &boxID, + &pts, + &messageDate, + ) + if errors.Is(err, pgx.ErrNoRows) { + return loginCodeDeliveryReceipt{}, false, nil + } + if err != nil { + return loginCodeDeliveryReceipt{}, false, fmt.Errorf("load login code delivery receipt: %w", err) + } + receipt.messageBoxID = int(boxID) + receipt.pts = int(pts) + receipt.messageDate = int(messageDate) + return receipt, true, nil +} + +func (s *MessageStore) nextLoginCodeBoxID(ctx context.Context, qtx *sqlcgen.Queries, userID int64) (int, error) { + // The default allocator queries PostgreSQL. Run that query on the active + // transaction connection: querying s.q while holding the transaction can + // deadlock a MaxConns=1 pool. External allocators (Redis/counters) retain + // their normal semantics. + switch s.boxIDs.(type) { + case pgBoxIDAllocator, *pgBoxIDAllocator: + current, err := qtx.MaxMessageBoxID(ctx, userID) + if err != nil { + return 0, err + } + return int(current) + 1, nil + default: + return s.boxIDs.NextBoxID(ctx, userID) + } +} diff --git a/internal/store/postgres/login_code_delivery_integration_test.go b/internal/store/postgres/login_code_delivery_integration_test.go new file mode 100644 index 00000000..5d48c2a9 --- /dev/null +++ b/internal/store/postgres/login_code_delivery_integration_test.go @@ -0,0 +1,475 @@ +package postgres + +import ( + "context" + "errors" + "fmt" + "reflect" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "telesrv/internal/domain" + "telesrv/internal/store" +) + +func TestLoginCodeDeliveryPostgresAtomicFactsAndReplay(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + user := createLoginCodeDeliveryTestUser(t, ctx, pool, "basic") + req := domain.LoginCodeDeliveryRequest{ + UserID: user.ID, + PhoneCodeHash: "pg-login-code-basic-" + randomSuffix(t), + Code: "12345", + Date: 1700001000, + ExpiresAt: 1700001300, + } + + first, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, req) + if err != nil { + t.Fatalf("DeliverLoginCodeMessage: %v", err) + } + if !first.Created || first.Message.ID != 1 || first.Message.Pts != 1 || first.Message.UID <= 0 || first.Message.Out || + first.Message.OwnerUserID != user.ID || first.Message.Peer.ID != domain.OfficialSystemUserID || first.Message.From.ID != domain.OfficialSystemUserID { + t.Fatalf("first delivery = %+v, want first incoming 777000 message", first) + } + + assertLoginCodeDeliveryFacts(t, ctx, pool, user.ID, first.Message, 1) + var senderUserID, recipientUserID, randomID int64 + var delivered bool + var senderBoxID, senderPts, recipientBoxID, recipientPts int32 + if err := pool.QueryRow(ctx, ` +SELECT sender_user_id, + recipient_user_id, + random_id, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts +FROM private_messages +WHERE sender_user_id = $1 AND id = $2`, domain.OfficialSystemUserID, first.Message.UID).Scan( + &senderUserID, + &recipientUserID, + &randomID, + &delivered, + &senderBoxID, + &senderPts, + &recipientBoxID, + &recipientPts, + ); err != nil { + t.Fatalf("load private message receipt: %v", err) + } + if senderUserID != domain.OfficialSystemUserID || recipientUserID != user.ID || randomID != 0 || !delivered || + senderBoxID != 0 || senderPts != 0 || int(recipientBoxID) != first.Message.ID || int(recipientPts) != first.Message.Pts { + t.Fatalf("private receipt sender=%d recipient=%d random=%d delivered=%v sender=%d/%d recipient=%d/%d", + senderUserID, recipientUserID, randomID, delivered, senderBoxID, senderPts, recipientBoxID, recipientPts) + } + var officialSenderBoxes int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM message_boxes WHERE owner_user_id = $1 AND private_message_id = $2`, domain.OfficialSystemUserID, first.Message.UID).Scan(&officialSenderBoxes); err != nil { + t.Fatalf("count official sender boxes: %v", err) + } + if officialSenderBoxes != 0 { + t.Fatalf("official sender boxes = %d, want recipient-only login notification", officialSenderBoxes) + } + + var deliveryKey, codeFingerprint []byte + if err := pool.QueryRow(ctx, ` +SELECT delivery_key, code_fingerprint +FROM login_code_message_deliveries +WHERE user_id = $1 AND message_box_id = $2`, user.ID, first.Message.ID).Scan(&deliveryKey, &codeFingerprint); err != nil { + t.Fatalf("load compact receipt: %v", err) + } + if len(deliveryKey) != 32 || len(codeFingerprint) != 32 || string(deliveryKey) == req.PhoneCodeHash { + t.Fatalf("compact receipt key/fingerprint lengths = %d/%d", len(deliveryKey), len(codeFingerprint)) + } + + replayReq := req + replayReq.Date += 99 + replay, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, replayReq) + if err != nil { + t.Fatalf("replay DeliverLoginCodeMessage: %v", err) + } + if replay.Created || !reflect.DeepEqual(replay.Message, first.Message) { + t.Fatalf("replay = %+v, want immutable first result %+v", replay, first) + } + assertLoginCodeDeliveryFacts(t, ctx, pool, user.ID, first.Message, 1) + + changedCode := req + changedCode.Code = "54321" + if _, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, changedCode); !errors.Is(err, domain.ErrLoginCodeDeliveryConflict) { + t.Fatalf("changed-code replay err = %v, want ErrLoginCodeDeliveryConflict", err) + } + assertLoginCodeDeliveryFacts(t, ctx, pool, user.ID, first.Message, 1) + + otherUser := createLoginCodeDeliveryTestUser(t, ctx, pool, "conflict") + changedUser := req + changedUser.UserID = otherUser.ID + if _, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, changedUser); !errors.Is(err, domain.ErrLoginCodeDeliveryConflict) { + t.Fatalf("changed-user replay err = %v, want ErrLoginCodeDeliveryConflict", err) + } + var otherFacts int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM message_boxes WHERE owner_user_id = $1`, otherUser.ID).Scan(&otherFacts); err != nil { + t.Fatalf("count changed-user facts: %v", err) + } + if otherFacts != 0 { + t.Fatalf("changed-user replay created %d message boxes", otherFacts) + } +} + +func TestLoginCodeDeliveryPostgresConcurrentExactlyOnce(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + user := createLoginCodeDeliveryTestUser(t, ctx, pool, "concurrent") + req := domain.LoginCodeDeliveryRequest{ + UserID: user.ID, + PhoneCodeHash: "pg-login-code-concurrent-" + randomSuffix(t), + Code: "24680", + Date: 1700001100, + ExpiresAt: 1700001400, + } + + const workers = 24 + var created atomic.Int32 + results := make(chan domain.LoginCodeDeliveryResult, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + got, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, req) + if err != nil { + errs <- err + return + } + if got.Created { + created.Add(1) + } + results <- got + }() + } + wg.Wait() + close(errs) + close(results) + for err := range errs { + t.Fatalf("concurrent delivery: %v", err) + } + if created.Load() != 1 { + t.Fatalf("created calls = %d, want exactly 1", created.Load()) + } + var first domain.Message + for got := range results { + if first.ID == 0 { + first = got.Message + continue + } + if !reflect.DeepEqual(got.Message, first) { + t.Fatalf("concurrent result = %+v, want %+v", got.Message, first) + } + } + if first.ID == 0 { + t.Fatal("no successful concurrent result") + } + assertLoginCodeDeliveryFacts(t, ctx, pool, user.ID, first, 1) +} + +func TestLoginCodeDeliveryPostgresCommitAckLossRecoversFromReceipt(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + user := createLoginCodeDeliveryTestUser(t, ctx, pool, "commit-ack-loss") + req := domain.LoginCodeDeliveryRequest{ + UserID: user.ID, + PhoneCodeHash: "pg-login-code-commit-ack-loss-" + randomSuffix(t), + Code: "86420", + Date: int(time.Now().Unix()), + ExpiresAt: time.Now().Add(5 * time.Minute).Unix(), + } + + got, err := NewMessageStore(&commitAckLossDB{Pool: pool}).DeliverLoginCodeMessage(ctx, req) + if err != nil { + t.Fatalf("DeliverLoginCodeMessage with lost commit ACK: %v", err) + } + if got.Created { + t.Fatalf("commit-ACK recovery Created = true, want conservative replay result") + } + assertLoginCodeDeliveryFacts(t, ctx, pool, user.ID, got.Message, 1) + + replay, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, req) + if err != nil { + t.Fatalf("replay after lost commit ACK: %v", err) + } + if replay.Created || !reflect.DeepEqual(replay.Message, got.Message) { + t.Fatalf("replay = %+v, want recovered snapshot %+v", replay, got) + } +} + +func TestLoginCodeDeliveryPostgresDifferentUsersDoNotRewriteOfficialUser(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + firstUser := createLoginCodeDeliveryTestUser(t, ctx, pool, "official-row-first") + now := int(time.Now().Unix()) + if _, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{ + UserID: firstUser.ID, PhoneCodeHash: "official-row-first-" + randomSuffix(t), Code: "12345", Date: now, ExpiresAt: int64(now + 300), + }); err != nil { + t.Fatalf("first delivery: %v", err) + } + var xminBefore string + if err := pool.QueryRow(ctx, `SELECT xmin::text FROM users WHERE id = $1`, domain.OfficialSystemUserID).Scan(&xminBefore); err != nil { + t.Fatalf("load official user xmin: %v", err) + } + + const workers = 12 + users := make([]domain.User, workers) + hashes := make([]string, workers) + for i := range users { + users[i] = createLoginCodeDeliveryTestUser(t, ctx, pool, fmt.Sprintf("official-row-%02d", i)) + hashes[i] = fmt.Sprintf("official-row-concurrent-%d-%s", i, randomSuffix(t)) + } + var wg sync.WaitGroup + errs := make(chan error, workers) + for i, user := range users { + wg.Add(1) + go func(i int, user domain.User) { + defer wg.Done() + _, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{ + UserID: user.ID, PhoneCodeHash: hashes[i], Code: "12345", Date: now, ExpiresAt: int64(now + 300), + }) + if err != nil { + errs <- err + } + }(i, user) + } + wg.Wait() + close(errs) + for err := range errs { + t.Fatalf("different-user delivery: %v", err) + } + var xminAfter string + if err := pool.QueryRow(ctx, `SELECT xmin::text FROM users WHERE id = $1`, domain.OfficialSystemUserID).Scan(&xminAfter); err != nil { + t.Fatalf("reload official user xmin: %v", err) + } + if xminAfter != xminBefore { + t.Fatalf("official system user row was rewritten: xmin %s -> %s", xminBefore, xminAfter) + } +} + +func TestLoginCodeDeliveryPostgresReceiptRetentionIsBoundedAndSeekOrdered(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + now := time.Now().Truncate(time.Second) + users := []domain.User{ + createLoginCodeDeliveryTestUser(t, ctx, pool, "expiry-old-1"), + createLoginCodeDeliveryTestUser(t, ctx, pool, "expiry-old-2"), + createLoginCodeDeliveryTestUser(t, ctx, pool, "expiry-future"), + } + for i, user := range users { + expiresAt := now.Add(time.Hour).Unix() + if i < 2 { + expiresAt = now.Add(time.Duration(i-2) * time.Minute).Unix() + } + if _, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{ + UserID: user.ID, PhoneCodeHash: fmt.Sprintf("expiry-%d-%s", i, randomSuffix(t)), Code: "12345", Date: int(now.Add(-time.Hour).Unix()), ExpiresAt: expiresAt, + }); err != nil { + t.Fatalf("seed expiry receipt %d: %v", i, err) + } + } + store := NewMessageStore(pool) + deleted, err := store.DeleteExpiredLoginCodeDeliveries(ctx, now, 1) + if err != nil || deleted != 1 { + t.Fatalf("first bounded retention = %d, %v; want 1", deleted, err) + } + deleted, err = store.DeleteExpiredLoginCodeDeliveries(ctx, now, 10) + if err != nil || deleted != 1 { + t.Fatalf("second bounded retention = %d, %v; want 1", deleted, err) + } + var receipts, messages, events int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM login_code_message_deliveries WHERE user_id = ANY($1)`, []int64{users[0].ID, users[1].ID, users[2].ID}).Scan(&receipts); err != nil { + t.Fatalf("count retained receipts: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM message_boxes WHERE owner_user_id = ANY($1) AND peer_id = $2`, []int64{users[0].ID, users[1].ID, users[2].ID}, domain.OfficialSystemUserID).Scan(&messages); err != nil { + t.Fatalf("count retained messages: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM user_update_events WHERE user_id = ANY($1) AND event_type = 'new_message'`, []int64{users[0].ID, users[1].ID, users[2].ID}).Scan(&events); err != nil { + t.Fatalf("count retained events: %v", err) + } + if receipts != 1 || messages != 3 || events != 3 { + t.Fatalf("after receipt GC receipts/messages/events = %d/%d/%d, want 1/3/3", receipts, messages, events) + } +} + +func TestLoginCodeDeliveryPostgresRollsBackEveryFact(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + user := createLoginCodeDeliveryTestUser(t, ctx, pool, "rollback") + first, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{ + UserID: user.ID, + PhoneCodeHash: "pg-login-code-rollback-first-" + randomSuffix(t), + Code: "11111", + Date: 1700001200, + ExpiresAt: 1700001500, + }) + if err != nil { + t.Fatalf("seed first delivery: %v", err) + } + if first.Message.ID != 1 || first.Message.Pts != 1 { + t.Fatalf("first allocation = id %d pts %d, want 1/1", first.Message.ID, first.Message.Pts) + } + + failedReq := domain.LoginCodeDeliveryRequest{ + UserID: user.ID, + PhoneCodeHash: "pg-login-code-rollback-failed-" + randomSuffix(t), + Code: "22222", + Date: 1700001201, + ExpiresAt: 1700001501, + } + failing := NewMessageStore(pool, WithMessageAllocators(loginCodeFixedBoxAllocator{boxID: first.Message.ID})) + if _, err := failing.DeliverLoginCodeMessage(ctx, failedReq); err == nil { + t.Fatal("duplicate box allocator delivery succeeded, want rollback") + } + assertLoginCodeDeliveryFacts(t, ctx, pool, user.ID, first.Message, 1) + failedKey, err := store.LoginCodeDeliveryKey(failedReq.PhoneCodeHash) + if err != nil { + t.Fatalf("failed delivery key: %v", err) + } + var failedReceipts, failedBodies int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM login_code_message_deliveries WHERE delivery_key = $1`, failedKey[:]).Scan(&failedReceipts); err != nil { + t.Fatalf("count failed receipts: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM private_messages WHERE sender_user_id = $1 AND recipient_user_id = $2 AND body LIKE '%22222%'`, domain.OfficialSystemUserID, user.ID).Scan(&failedBodies); err != nil { + t.Fatalf("count failed private messages: %v", err) + } + if failedReceipts != 0 || failedBodies != 0 { + t.Fatalf("failed transaction leaked receipts=%d private_messages=%d", failedReceipts, failedBodies) + } + + third, err := NewMessageStore(pool).DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{ + UserID: user.ID, + PhoneCodeHash: "pg-login-code-rollback-third-" + randomSuffix(t), + Code: "33333", + Date: 1700001202, + ExpiresAt: 1700001502, + }) + if err != nil { + t.Fatalf("delivery after rollback: %v", err) + } + if third.Message.ID != 2 || third.Message.Pts != 2 { + t.Fatalf("allocation after rollback = id %d pts %d, want contiguous 2/2", third.Message.ID, third.Message.Pts) + } +} + +type loginCodeFixedBoxAllocator struct { + boxID int +} + +type commitAckLossDB struct { + *pgxpool.Pool +} + +func (d *commitAckLossDB) Begin(ctx context.Context) (pgx.Tx, error) { + tx, err := d.Pool.Begin(ctx) + if err != nil { + return nil, err + } + return &commitAckLossTx{Tx: tx}, nil +} + +type commitAckLossTx struct { + pgx.Tx +} + +func (t *commitAckLossTx) Commit(ctx context.Context) error { + if err := t.Tx.Commit(ctx); err != nil { + return err + } + return errors.New("synthetic lost commit acknowledgement") +} + +func (a loginCodeFixedBoxAllocator) NextBoxID(context.Context, int64) (int, error) { + return a.boxID, nil +} + +func (a loginCodeFixedBoxAllocator) CurrentBoxID(context.Context, int64) (int, error) { + return a.boxID, nil +} + +func createLoginCodeDeliveryTestUser(t *testing.T, ctx context.Context, pool *pgxpool.Pool, label string) domain.User { + t.Helper() + user, err := NewUserStore(pool).Create(ctx, domain.User{ + AccessHash: 8100000000, + Phone: "+1888" + randomSuffix(t), + FirstName: "LoginCode" + label, + }) + if err != nil { + t.Fatalf("create login code test user: %v", err) + } + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), `DELETE FROM users WHERE id = $1`, user.ID) + }) + return user +} + +func assertLoginCodeDeliveryFacts(t *testing.T, ctx context.Context, pool *pgxpool.Pool, userID int64, msg domain.Message, want int) { + t.Helper() + queries := []struct { + name string + sql string + args []any + }{ + {"private_messages", `SELECT count(*) FROM private_messages WHERE sender_user_id = $1 AND recipient_user_id = $2`, []any{domain.OfficialSystemUserID, userID}}, + {"message_boxes", `SELECT count(*) FROM message_boxes WHERE owner_user_id = $1 AND peer_type = 'user' AND peer_id = $2`, []any{userID, domain.OfficialSystemUserID}}, + {"dialogs", `SELECT count(*) FROM dialogs WHERE user_id = $1 AND peer_type = 'user' AND peer_id = $2`, []any{userID, domain.OfficialSystemUserID}}, + {"user_update_events", `SELECT count(*) FROM user_update_events WHERE user_id = $1 AND event_type = 'new_message'`, []any{userID}}, + {"dispatch_outbox", `SELECT count(*) FROM dispatch_outbox WHERE target_user_id = $1 AND event_type = 'new_message'`, []any{userID}}, + {"delivery_receipts", `SELECT count(*) FROM login_code_message_deliveries WHERE user_id = $1`, []any{userID}}, + } + for _, query := range queries { + var got int + if err := pool.QueryRow(ctx, query.sql, query.args...).Scan(&got); err != nil { + t.Fatalf("count %s: %v", query.name, err) + } + if got != want { + t.Fatalf("%s count = %d, want %d", query.name, got, want) + } + } + + var boxPts, eventPts, eventBoxID, outboxPts int32 + var eventType, outboxEventType string + if err := pool.QueryRow(ctx, ` +SELECT b.pts, + e.pts, + e.message_box_id, + e.event_type, + o.pts, + o.event_type +FROM message_boxes b +JOIN user_update_events e + ON e.user_id = b.owner_user_id + AND e.message_box_id = b.box_id +JOIN dispatch_outbox o + ON o.target_user_id = e.user_id + AND o.pts = e.pts +WHERE b.owner_user_id = $1 + AND b.box_id = $2`, userID, msg.ID).Scan(&boxPts, &eventPts, &eventBoxID, &eventType, &outboxPts, &outboxEventType); err != nil { + t.Fatalf("load login message/event/outbox chain: %v", err) + } + if int(boxPts) != msg.Pts || eventPts != boxPts || eventBoxID != int32(msg.ID) || eventType != string(domain.UpdateEventNewMessage) || + outboxPts != eventPts || outboxEventType != eventType { + t.Fatalf("box/event/outbox chain = box_pts %d event %d/%d/%s outbox %d/%s, message=%+v", + boxPts, eventPts, eventBoxID, eventType, outboxPts, outboxEventType, msg) + } + var topMessageID, unreadCount int32 + if err := pool.QueryRow(ctx, ` +SELECT top_message_id, unread_count +FROM dialogs +WHERE user_id = $1 AND peer_type = 'user' AND peer_id = $2`, userID, domain.OfficialSystemUserID).Scan(&topMessageID, &unreadCount); err != nil { + t.Fatalf("load login code dialog: %v", err) + } + if int(topMessageID) != msg.ID || int(unreadCount) != want { + t.Fatalf("dialog top/unread = %d/%d, want %d/%d", topMessageID, unreadCount, msg.ID, want) + } +} diff --git a/internal/store/postgres/login_code_delivery_retention.go b/internal/store/postgres/login_code_delivery_retention.go new file mode 100644 index 00000000..a3d94bf4 --- /dev/null +++ b/internal/store/postgres/login_code_delivery_retention.go @@ -0,0 +1,32 @@ +package postgres + +import ( + "context" + "fmt" + "time" +) + +// DeleteExpiredLoginCodeDeliveries seek-deletes compact idempotency receipts +// whose corresponding opaque codes are no longer usable. Message/update facts +// are deliberately retained; only the replay receipt is ephemeral. +func (s *MessageStore) DeleteExpiredLoginCodeDeliveries(ctx context.Context, expiredBefore time.Time, limit int) (int, error) { + if limit <= 0 { + return 0, nil + } + tag, err := s.db.Exec(ctx, ` +WITH doomed AS ( + SELECT delivery_key + FROM login_code_message_deliveries + WHERE expires_at <= $1 + ORDER BY expires_at, delivery_key + LIMIT $2 + FOR UPDATE SKIP LOCKED +) +DELETE FROM login_code_message_deliveries AS d +USING doomed +WHERE d.delivery_key = doomed.delivery_key`, expiredBefore.UTC(), limit) + if err != nil { + return 0, fmt.Errorf("delete expired login code delivery receipts: %w", err) + } + return int(tag.RowsAffected()), nil +} diff --git a/internal/store/postgres/media.go b/internal/store/postgres/media.go index 0b5dcb16..52308aa5 100644 --- a/internal/store/postgres/media.go +++ b/internal/store/postgres/media.go @@ -44,6 +44,69 @@ func bytesOrEmpty(b []byte) []byte { var _ store.MediaStore = (*MediaStore)(nil) +func validUploadedMediaReceipt(receipt domain.UploadedMediaReceipt) bool { + if receipt.OwnerUserID == 0 || receipt.FileID == 0 || receipt.MediaID == 0 || len(receipt.IntentHash) != 32 { + return false + } + return receipt.Kind == domain.UploadedMediaPhoto || receipt.Kind == domain.UploadedMediaDocument +} + +func (s *MediaStore) GetUploadedMediaReceipt(ctx context.Context, ownerUserID, fileID int64) (domain.UploadedMediaReceipt, bool, error) { + var receipt domain.UploadedMediaReceipt + var kind string + err := s.db.QueryRow(ctx, ` +SELECT owner_user_id, file_id, intent_hash, media_kind, media_id, created_at +FROM uploaded_media_receipts +WHERE owner_user_id = $1 AND file_id = $2`, ownerUserID, fileID).Scan( + &receipt.OwnerUserID, + &receipt.FileID, + &receipt.IntentHash, + &kind, + &receipt.MediaID, + &receipt.CreatedAt, + ) + if errors.Is(err, pgx.ErrNoRows) { + return domain.UploadedMediaReceipt{}, false, nil + } + if err != nil { + return domain.UploadedMediaReceipt{}, false, fmt.Errorf("get uploaded media receipt: %w", err) + } + receipt.Kind = domain.UploadedMediaKind(kind) + if !validUploadedMediaReceipt(receipt) { + return domain.UploadedMediaReceipt{}, false, fmt.Errorf( + "get uploaded media receipt: invalid owner=%d file=%d kind=%q media=%d hash=%d", + receipt.OwnerUserID, receipt.FileID, receipt.Kind, receipt.MediaID, len(receipt.IntentHash), + ) + } + return receipt, true, nil +} + +func (s *MediaStore) PutUploadedMediaReceipt(ctx context.Context, receipt domain.UploadedMediaReceipt) (domain.UploadedMediaReceipt, bool, error) { + if !validUploadedMediaReceipt(receipt) { + return domain.UploadedMediaReceipt{}, false, fmt.Errorf( + "put uploaded media receipt: invalid owner=%d file=%d kind=%q media=%d hash=%d", + receipt.OwnerUserID, receipt.FileID, receipt.Kind, receipt.MediaID, len(receipt.IntentHash), + ) + } + tag, err := s.db.Exec(ctx, ` +INSERT INTO uploaded_media_receipts (owner_user_id, file_id, intent_hash, media_kind, media_id) +VALUES ($1, $2, $3, $4, $5) +ON CONFLICT (owner_user_id, file_id) DO NOTHING`, + receipt.OwnerUserID, receipt.FileID, receipt.IntentHash, string(receipt.Kind), receipt.MediaID, + ) + if err != nil { + return domain.UploadedMediaReceipt{}, false, fmt.Errorf("put uploaded media receipt: %w", err) + } + stored, found, err := s.GetUploadedMediaReceipt(ctx, receipt.OwnerUserID, receipt.FileID) + if err != nil { + return domain.UploadedMediaReceipt{}, false, err + } + if !found { + return domain.UploadedMediaReceipt{}, false, fmt.Errorf("put uploaded media receipt: row disappeared after insert") + } + return stored, tag.RowsAffected() == 1, nil +} + // ---- 上传分片 ---- func (s *MediaStore) SaveFilePart(ctx context.Context, part domain.UploadPart) error { diff --git a/internal/store/postgres/media_integration_test.go b/internal/store/postgres/media_integration_test.go index 717fda41..39e0b650 100644 --- a/internal/store/postgres/media_integration_test.go +++ b/internal/store/postgres/media_integration_test.go @@ -220,6 +220,43 @@ func TestMediaStoreRoundTrip(t *testing.T) { } } +func TestUploadedMediaReceiptFirstWriterWinsPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + media := NewMediaStore(pool) + userID := createRevokeTestUser(t, ctx, pool, "upload-receipt") + const fileID = int64(880055501) + t.Cleanup(func() { + _, _ = pool.Exec(context.Background(), "DELETE FROM uploaded_media_receipts WHERE owner_user_id = $1 AND file_id = $2", userID, fileID) + }) + first := domain.UploadedMediaReceipt{ + OwnerUserID: userID, + FileID: fileID, + IntentHash: bytes.Repeat([]byte{1}, 32), + Kind: domain.UploadedMediaPhoto, + MediaID: 7001, + } + stored, created, err := media.PutUploadedMediaReceipt(ctx, first) + if err != nil || !created || stored.MediaID != first.MediaID || !bytes.Equal(stored.IntentHash, first.IntentHash) { + t.Fatalf("first receipt = %+v created=%v err=%v", stored, created, err) + } + second := first + second.IntentHash = bytes.Repeat([]byte{2}, 32) + second.Kind = domain.UploadedMediaDocument + second.MediaID = 7002 + stored, created, err = media.PutUploadedMediaReceipt(ctx, second) + if err != nil || created { + t.Fatalf("conflicting receipt created=%v err=%v", created, err) + } + if stored.Kind != first.Kind || stored.MediaID != first.MediaID || !bytes.Equal(stored.IntentHash, first.IntentHash) { + t.Fatalf("conflicting receipt replaced first writer: %+v", stored) + } + got, found, err := media.GetUploadedMediaReceipt(ctx, userID, fileID) + if err != nil || !found || got.MediaID != first.MediaID { + t.Fatalf("get receipt = %+v found=%v err=%v", got, found, err) + } +} + func TestMediaStoreDocumentCacheCopiesAndRefreshes(t *testing.T) { pool := testPool(t) ctx := context.Background() diff --git a/internal/store/postgres/message_delete.go b/internal/store/postgres/message_delete.go index fc20dfba..ab69c730 100644 --- a/internal/store/postgres/message_delete.go +++ b/internal/store/postgres/message_delete.go @@ -195,6 +195,29 @@ func (s *MessageStore) finishDeleteMessagesTx(ctx context.Context, db sqlcgen.DB Date: date, MessageIDs: ids, } + deleteIDsJSON, err := encodeEventMessageIDs(event.MessageIDs) + if err != nil { + return res, fmt.Errorf("encode sender delete receipt ids: %w", err) + } + senderPrivateIDs := make([]int64, 0, len(rows)) + for _, row := range rows { + if row.ownerUserID == userID && row.messageSenderID == userID && row.privateMessageID != 0 { + senderPrivateIDs = append(senderPrivateIDs, row.privateMessageID) + } + } + if len(senderPrivateIDs) > 0 { + if _, err := db.Exec(ctx, ` +UPDATE private_messages +SET sender_delete_pts = $3, + sender_delete_pts_count = $4, + sender_delete_date = $5, + sender_delete_message_ids = $6::jsonb +WHERE sender_user_id = $1 + AND id = ANY($2::bigint[]) + AND sender_box_id > 0`, userID, senderPrivateIDs, event.Pts, event.PtsCount, event.Date, deleteIDsJSON); err != nil { + return res, fmt.Errorf("save sender delete replay receipt: %w", err) + } + } if err := appendDeleteMessagesEvent(ctx, q, event); err != nil { return res, err } diff --git a/internal/store/postgres/message_delete_integration_test.go b/internal/store/postgres/message_delete_integration_test.go index ce506baf..ec47945b 100644 --- a/internal/store/postgres/message_delete_integration_test.go +++ b/internal/store/postgres/message_delete_integration_test.go @@ -490,6 +490,12 @@ func TestMessageStoreDeleteHistoryBatchesHugeMaxID(t *testing.T) { sender_user_id, recipient_user_id, random_id, + request_fingerprint, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts, message_date, body, entities @@ -498,6 +504,12 @@ func TestMessageStoreDeleteHistoryBatchesHugeMaxID(t *testing.T) { $1::bigint, $2::bigint, 910000000 + g, + decode(repeat('00', 32), 'hex'), + false, + g, + g, + 0, + 0, 1700002000 + g, 'bulk history', '[]'::jsonb @@ -530,7 +542,7 @@ func TestMessageStoreDeleteHistoryBatchesHugeMaxID(t *testing.T) { true, 'bulk history', '[]'::jsonb, - 0 + (random_id - 910000000)::int FROM pm `, owner.ID, peerUser.ID, total); err != nil { t.Fatalf("seed bulk history: %v", err) diff --git a/internal/store/postgres/message_markup_integration_test.go b/internal/store/postgres/message_markup_integration_test.go index 1f5abd0c..173f69ee 100644 --- a/internal/store/postgres/message_markup_integration_test.go +++ b/internal/store/postgres/message_markup_integration_test.go @@ -3,6 +3,7 @@ package postgres import ( "bytes" "context" + "errors" "testing" "time" @@ -193,15 +194,30 @@ func TestSendPrivateViaBotIDSurvivesReadPaths(t *testing.T) { SenderUserID: sender.ID, RecipientUserID: recipient.ID, RandomID: randomID, - Message: "inline via duplicate", + Message: "inline via", ViaBotID: viaBotID, Date: int(time.Now().Unix()), }) if err != nil { t.Fatalf("duplicate private via bot: %v", err) } - if !dup.Duplicate || dup.SenderMessage.ViaBotID != viaBotID || dup.RecipientMessage.ViaBotID != viaBotID { - t.Fatalf("duplicate via = duplicate %v sender %d recipient %d, want duplicate true via %d", dup.Duplicate, dup.SenderMessage.ViaBotID, dup.RecipientMessage.ViaBotID, viaBotID) + if !dup.Duplicate || + dup.SenderMessage.ID != res.SenderMessage.ID || dup.SenderMessage.Pts != res.SenderMessage.Pts || + dup.RecipientMessage.ID != res.RecipientMessage.ID || dup.RecipientMessage.Pts != res.RecipientMessage.Pts { + t.Fatalf("duplicate immutable receipt = %+v/%+v, want sender id/pts %d/%d recipient %d/%d", + dup.SenderMessage, dup.RecipientMessage, + res.SenderMessage.ID, res.SenderMessage.Pts, + res.RecipientMessage.ID, res.RecipientMessage.Pts) + } + if _, err := messages.SendPrivateText(ctx, domain.SendPrivateTextRequest{ + SenderUserID: sender.ID, + RecipientUserID: recipient.ID, + RandomID: randomID, + Message: "inline via conflict", + ViaBotID: viaBotID, + Date: int(time.Now().Unix()), + }); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("conflicting via-bot replay err = %v, want ErrMessageRandomIDDuplicate", err) } recipientHistory, err := messages.ListByUser(ctx, recipient.ID, domain.MessageFilter{Limit: 10}) diff --git a/internal/store/postgres/message_plan_integration_test.go b/internal/store/postgres/message_plan_integration_test.go index 954570e6..2853f038 100644 --- a/internal/store/postgres/message_plan_integration_test.go +++ b/internal/store/postgres/message_plan_integration_test.go @@ -194,44 +194,44 @@ ORDER BY owner_user_id ASC, box_id ASC requireUniquePartitionCountAtMost(t, deleteByPrivatePlan, `message_boxes_p\d+`, 2) dispatchPlan := explainText(t, ctx, tx, ` -WITH picked AS ( - SELECT target_user_id, pts, id - FROM dispatch_outbox - WHERE ( - status = 'pending' - AND next_attempt_at <= now() +WITH picked_heads AS ( + SELECT h.target_user_id, h.head_id, h.head_pts + FROM dispatch_outbox_user_heads h + WHERE h.logical_shard = ANY($1::smallint[]) + AND ( + (h.status = 'pending' AND h.next_attempt_at <= now()) + OR + (h.status = 'dispatching' AND h.updated_at < now() - interval '30 seconds') ) - OR ( - status = 'dispatching' - AND updated_at < now() - interval '30 seconds' - ) - ORDER BY next_attempt_at ASC, target_user_id ASC, id ASC + ORDER BY h.next_attempt_at ASC, h.target_user_id ASC, h.head_pts ASC, h.head_id ASC LIMIT 100 - FOR UPDATE SKIP LOCKED + FOR UPDATE OF h SKIP LOCKED ) -SELECT target_user_id, pts, id -FROM picked -`) +SELECT d.target_user_id, d.pts, d.id +FROM picked_heads h +JOIN dispatch_outbox d + ON d.target_user_id = h.target_user_id + AND d.id = h.head_id +`, []int16{int16(recipient.ID % 256)}) + requirePlanContains(t, dispatchPlan, "dispatch_outbox_user_heads_dispatching_shard_idx") requirePlanContains(t, dispatchPlan, "dispatch_outbox") requirePlanContains(t, dispatchPlan, "Index") requirePlanNotMatches(t, dispatchPlan, `dispatch_outbox_p\d+`) requirePlanNotContains(t, dispatchPlan, "Seq Scan") failedCleanupPlan := explainText(t, ctx, tx, ` -WITH doomed AS ( - SELECT target_user_id, id - FROM dispatch_outbox +WITH doomed AS MATERIALIZED ( + SELECT target_user_id, head_id AS id + FROM dispatch_outbox_user_heads WHERE status = 'failed' - AND updated_at < now() - interval '1 day' - ORDER BY updated_at ASC, target_user_id ASC, id ASC + AND updated_at < now() - interval '1 minute' + ORDER BY updated_at ASC, target_user_id ASC, head_id ASC LIMIT 100 ) SELECT target_user_id, id FROM doomed `) - requirePlanContains(t, failedCleanupPlan, "dispatch_outbox") - requirePlanContains(t, failedCleanupPlan, "Index") - requirePlanNotMatches(t, failedCleanupPlan, `dispatch_outbox_p\d+`) + requirePlanContains(t, failedCleanupPlan, "dispatch_outbox_user_heads_failed_cleanup_idx") requirePlanNotContains(t, failedCleanupPlan, "Seq Scan") } diff --git a/internal/store/postgres/message_send.go b/internal/store/postgres/message_send.go index 62f0ed64..306e5cce 100644 --- a/internal/store/postgres/message_send.go +++ b/internal/store/postgres/message_send.go @@ -10,6 +10,7 @@ import ( "github.com/jackc/pgx/v5/pgconn" "sort" "telesrv/internal/domain" + "telesrv/internal/store" "telesrv/internal/store/postgres/sqlcgen" "time" ) @@ -50,6 +51,10 @@ func (s *MessageStore) Create(ctx context.Context, msg domain.Message) (domain.M } func (s *MessageStore) ensureOfficialSystemUser(ctx context.Context, msg domain.Message) error { + return ensureOfficialSystemUserWithDB(ctx, s.db, msg) +} + +func ensureOfficialSystemUserWithDB(ctx context.Context, db sqlcgen.DBTX, msg domain.Message) error { if msg.Peer.Type != domain.PeerTypeUser && msg.From.Type != domain.PeerTypeUser { return nil } @@ -60,22 +65,10 @@ func (s *MessageStore) ensureOfficialSystemUser(ctx context.Context, msg domain. if !ok { return nil } - if _, err := s.db.Exec(ctx, ` + if _, err := db.Exec(ctx, ` INSERT INTO users (id, access_hash, phone, first_name, last_name, username, country_code, verified, support, about, is_bot, bot_info_version) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) -ON CONFLICT (id) DO UPDATE SET - access_hash = EXCLUDED.access_hash, - phone = EXCLUDED.phone, - first_name = EXCLUDED.first_name, - last_name = EXCLUDED.last_name, - username = EXCLUDED.username, - country_code = EXCLUDED.country_code, - verified = EXCLUDED.verified, - support = EXCLUDED.support, - about = EXCLUDED.about, - is_bot = EXCLUDED.is_bot, - bot_info_version = EXCLUDED.bot_info_version, - updated_at = now() +ON CONFLICT (id) DO NOTHING `, u.ID, u.AccessHash, u.Phone, u.FirstName, u.LastName, u.Username, u.CountryCode, u.Verified, u.Support, u.About, u.Bot, u.BotInfoVersion); err != nil { return fmt.Errorf("ensure official system user: %w", err) } @@ -129,6 +122,21 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP if err != nil { return domain.SendPrivateTextResult{}, err } + requestFingerprint, err := store.PrivateSendFingerprint(req) + if err != nil { + return domain.SendPrivateTextResult{}, err + } + // 常见的 lost-response 重放在开事务和拿双方 advisory lock 之前直接返回； + // 并发首次请求仍由事务内 unique conflict + qtx 兜底，不能只依赖本次预查。 + // RPC/app 已完成同一只读查询时可跳过这次重复 round-trip。 + if !req.IdempotencyPreflighted { + if duplicate, found, err := s.duplicateSendResult(ctx, s.q, req, requestFingerprint); err != nil { + return domain.SendPrivateTextResult{}, err + } else if found { + duplicate.Duplicate = true + return duplicate, nil + } + } senderReply, recipientReply, err := s.resolvePrivateSendReply(ctx, req) if err != nil { return domain.SendPrivateTextResult{}, err @@ -186,30 +194,36 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP } privateArg := sqlcgen.CreatePrivateMessageParams{ - SenderUserID: req.SenderUserID, - RecipientUserID: req.RecipientUserID, - RandomID: req.RandomID, - MessageDate: int32(req.Date), - Body: req.Message, - TtlPeriod: int32(ttlPeriod), - ExpiresAt: int32(expiresAt), - EntitiesJson: entities, - MediaJson: mediaJSON, - ReplyMarkupJson: replyMarkupJSON, - RichMessageJson: richMessageJSON, - ViaBotID: req.ViaBotID, - GroupedID: req.GroupedID, - Effect: req.Effect, + SenderUserID: req.SenderUserID, + RecipientUserID: req.RecipientUserID, + RandomID: req.RandomID, + RequestFingerprint: requestFingerprint, + RecipientDelivered: deliverRecipient, + MessageDate: int32(req.Date), + Body: req.Message, + TtlPeriod: int32(ttlPeriod), + ExpiresAt: int32(expiresAt), + EntitiesJson: entities, + MediaJson: mediaJSON, + ReplyMarkupJson: replyMarkupJSON, + RichMessageJson: richMessageJSON, + ViaBotID: req.ViaBotID, + GroupedID: req.GroupedID, + Effect: req.Effect, } applyCreatePrivateMessageMetadata(&privateArg, senderMeta) pm, err := qtx.CreatePrivateMessage(ctx, privateArg) if err != nil { if errors.Is(err, pgx.ErrNoRows) { - // 幂等重复：返回原消息盒；此时还没有分配 pts，重复发送不应制造额外事件。 - dup, dupErr := s.duplicateSendResult(ctx, req.SenderUserID, req.RecipientUserID, req.RandomID) + // 预查与 INSERT 之间另一请求可能已提交。必须在当前 qtx 读取， + // 不能持事务连接/advisory lock 再从 s.q 申请第二条池连接。 + dup, found, dupErr := s.duplicateSendResult(ctx, qtx, req, requestFingerprint) if dupErr != nil { return domain.SendPrivateTextResult{}, dupErr } + if !found { + return domain.SendPrivateTextResult{}, fmt.Errorf("duplicate private message disappeared after unique conflict") + } dup.Duplicate = true return dup, nil } @@ -267,6 +281,7 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP return domain.SendPrivateTextResult{}, fmt.Errorf("create sender box: %w", err) } sender := messageFromBoxRow(senderRow) + sender.RandomID = req.RandomID // 共享媒体索引(0118):发送者侧 box 按媒体类别建索引(peer=收件人)。 if err := insertMessageBoxMediaIndexTx(ctx, tx, req.SenderUserID, req.RecipientUserID, int(senderBoxID), req.Date, req.Media, req.Entities); err != nil { return domain.SendPrivateTextResult{}, err @@ -328,6 +343,7 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP return domain.SendPrivateTextResult{}, fmt.Errorf("create recipient box: %w", err) } recipient = messageFromBoxRow(recipientRow) + recipient.RandomID = req.RandomID // 共享媒体索引(0118):收件人侧 box 按媒体类别建索引(peer=发送者)。 if err := insertMessageBoxMediaIndexTx(ctx, tx, req.RecipientUserID, req.SenderUserID, int(recipientBoxID), req.Date, req.Media, req.Entities); err != nil { return domain.SendPrivateTextResult{}, err @@ -355,6 +371,33 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP } } + receiptRecipientBoxID, receiptRecipientPts := recipientBoxID, recipientPts + if selfMessage { + receiptRecipientBoxID, receiptRecipientPts = sender.ID, sender.Pts + } + senderSnapshot, err := store.EncodePrivateSendSnapshot(sender) + if err != nil { + return domain.SendPrivateTextResult{}, err + } + tag, err := tx.Exec(ctx, ` +UPDATE private_messages +SET sender_box_id = $3, + sender_pts = $4, + recipient_box_id = $5, + recipient_pts = $6, + sender_snapshot = $7::jsonb +WHERE sender_user_id = $1 + AND id = $2 + AND sender_box_id = 0 + AND sender_pts = 0 + AND sender_snapshot = '{}'::jsonb`, req.SenderUserID, pm.ID, sender.ID, sender.Pts, receiptRecipientBoxID, receiptRecipientPts, senderSnapshot) + if err != nil { + return domain.SendPrivateTextResult{}, fmt.Errorf("save private send receipt: %w", err) + } + if tag.RowsAffected() != 1 { + return domain.SendPrivateTextResult{}, fmt.Errorf("save private send receipt: private message %d already has or lost its immutable receipt", pm.ID) + } + if err := tx.Commit(ctx); err != nil { return domain.SendPrivateTextResult{}, fmt.Errorf("commit send message tx: %w", err) } @@ -367,6 +410,28 @@ func (s *MessageStore) sendPrivateTextOnce(ctx context.Context, req domain.SendP }, nil } +// LookupPrivateSendReplay reads an existing receipt without permission checks, source/media +// resolution, locks or allocations. The authenticated app/RPC layer supplies sender identity. +func (s *MessageStore) LookupPrivateSendReplay(ctx context.Context, lookup domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) { + if lookup.SenderUserID == 0 || lookup.RecipientUserID == 0 || lookup.RandomID == 0 { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("private send replay: invalid scope") + } + if err := store.ValidateSendFingerprint(lookup.IdempotencyFingerprint, "private send replay"); err != nil { + return domain.SendPrivateTextResult{}, false, err + } + res, found, err := s.duplicateSendResult(ctx, s.q, domain.SendPrivateTextRequest{ + SenderUserID: lookup.SenderUserID, + RecipientUserID: lookup.RecipientUserID, + RandomID: lookup.RandomID, + IdempotencyFingerprint: lookup.IdempotencyFingerprint, + }, lookup.IdempotencyFingerprint) + if err != nil || !found { + return domain.SendPrivateTextResult{}, found, err + } + res.Duplicate = true + return res, true, nil +} + type boxIDCounterBumper interface { BumpBoxIDAtLeast(ctx context.Context, userID int64, floor int) error } @@ -400,49 +465,96 @@ func isMessageBoxDuplicateKey(err error) bool { return pgErr.Code == "23505" && strings.Contains(pgErr.ConstraintName, "message_boxes") } -func (s *MessageStore) duplicateSendResult(ctx context.Context, senderUserID, recipientUserID, randomID int64) (domain.SendPrivateTextResult, error) { - pm, err := s.q.GetPrivateMessageByRandomID(ctx, sqlcgen.GetPrivateMessageByRandomIDParams{ - SenderUserID: senderUserID, - RandomID: randomID, +func (s *MessageStore) duplicateSendResult(ctx context.Context, q *sqlcgen.Queries, req domain.SendPrivateTextRequest, requestFingerprint []byte) (domain.SendPrivateTextResult, bool, error) { + pm, err := q.GetPrivateMessageByRandomID(ctx, sqlcgen.GetPrivateMessageByRandomIDParams{ + SenderUserID: req.SenderUserID, + RandomID: req.RandomID, }) if err != nil { - return domain.SendPrivateTextResult{}, fmt.Errorf("get duplicate private message: %w", err) + if errors.Is(err, pgx.ErrNoRows) { + return domain.SendPrivateTextResult{}, false, nil + } + return domain.SendPrivateTextResult{}, false, fmt.Errorf("get duplicate private message: %w", err) } - senderRow, err := s.q.GetMessageBoxByPrivateMessage(ctx, sqlcgen.GetMessageBoxByPrivateMessageParams{ - OwnerUserID: senderUserID, + if pm.SenderUserID != req.SenderUserID || + pm.RecipientUserID != req.RecipientUserID || + !store.SamePrivateSendFingerprint(pm.RequestFingerprint, requestFingerprint) { + return domain.SendPrivateTextResult{}, false, domain.ErrMessageRandomIDDuplicate + } + if pm.SenderBoxID <= 0 || pm.SenderPts <= 0 { + return domain.SendPrivateTextResult{}, false, fmt.Errorf( + "duplicate private message %d has invalid immutable sender receipt box=%d pts=%d", + pm.ID, pm.SenderBoxID, pm.SenderPts, + ) + } + firstSender, err := store.DecodePrivateSendSnapshot([]byte(pm.SenderSnapshotJson)) + if err != nil { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("decode duplicate private message %d sender snapshot: %w", pm.ID, err) + } + if firstSender.ID != int(pm.SenderBoxID) || firstSender.UID != pm.ID || firstSender.RandomID != pm.RandomID || + firstSender.OwnerUserID != pm.SenderUserID || firstSender.Pts != int(pm.SenderPts) { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("duplicate private message %d sender snapshot disagrees with immutable receipt", pm.ID) + } + sender := firstSender + currentRow, currentErr := q.GetMessageBoxByPrivateMessage(ctx, sqlcgen.GetMessageBoxByPrivateMessageParams{ + OwnerUserID: pm.SenderUserID, PrivateMessageID: pm.ID, }) - if err != nil { - return domain.SendPrivateTextResult{}, fmt.Errorf("get duplicate sender box: %w", err) + if currentErr == nil { + sender = messageFromGetBoxRow(currentRow) + sender.RandomID = pm.RandomID + } else if !errors.Is(currentErr, pgx.ErrNoRows) { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("get current duplicate private message %d sender box: %w", pm.ID, currentErr) + } + var replayDelete *domain.UpdateEvent + if errors.Is(currentErr, pgx.ErrNoRows) { + messageIDs, decodeErr := decodeEventMessageIDs(pm.SenderDeleteMessageIdsJson) + if decodeErr != nil { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("decode duplicate private message %d delete ids: %w", pm.ID, decodeErr) + } + if pm.SenderDeletePts <= 0 || pm.SenderDeletePtsCount <= 0 || len(messageIDs) == 0 { + return domain.SendPrivateTextResult{}, false, fmt.Errorf("duplicate private message %d sender box is absent without a durable delete receipt", pm.ID) + } + event := domain.UpdateEvent{ + UserID: pm.SenderUserID, + Type: domain.UpdateEventDeleteMessages, + Pts: int(pm.SenderDeletePts), + PtsCount: int(pm.SenderDeletePtsCount), + Date: int(pm.SenderDeleteDate), + MessageIDs: messageIDs, + } + replayDelete = &event } - sender := messageFromGetBoxRow(senderRow) recipient := domain.Message{} - if recipientUserID == senderUserID { + if req.RecipientUserID == req.SenderUserID { recipient = sender } - if recipientUserID != senderUserID { - recipientRow, err := s.q.GetMessageBoxByPrivateMessage(ctx, sqlcgen.GetMessageBoxByPrivateMessageParams{ - OwnerUserID: recipientUserID, - PrivateMessageID: pm.ID, - }) - if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - return domain.SendPrivateTextResult{ - SenderMessage: sender, - SenderEvent: eventFromMessage(sender), - RecipientEvent: domain.UpdateEvent{}, - }, nil - } - return domain.SendPrivateTextResult{}, fmt.Errorf("get duplicate recipient box: %w", err) + if req.RecipientUserID != req.SenderUserID && pm.RecipientDelivered { + if pm.RecipientBoxID <= 0 || pm.RecipientPts <= 0 { + return domain.SendPrivateTextResult{}, false, fmt.Errorf( + "duplicate private message %d declares recipient delivery with invalid immutable receipt box=%d pts=%d", + pm.ID, pm.RecipientBoxID, pm.RecipientPts, + ) + } + recipient = domain.Message{ + ID: int(pm.RecipientBoxID), + UID: pm.ID, + RandomID: pm.RandomID, + OwnerUserID: pm.RecipientUserID, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: pm.SenderUserID}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: pm.SenderUserID}, + Date: int(pm.MessageDate), + Out: false, + Pts: int(pm.RecipientPts), } - recipient = messageFromGetBoxRow(recipientRow) } return domain.SendPrivateTextResult{ - SenderMessage: sender, - RecipientMessage: recipient, - SenderEvent: eventFromMessage(sender), - RecipientEvent: eventFromMessage(recipient), - }, nil + SenderMessage: sender, + RecipientMessage: recipient, + SenderEvent: eventFromMessage(firstSender), + RecipientEvent: eventFromMessage(recipient), + ReplayDeleteEvent: replayDelete, + }, true, nil } func (s *MessageStore) resolvePrivateSendReply(ctx context.Context, req domain.SendPrivateTextRequest) (*domain.MessageReply, *domain.MessageReply, error) { diff --git a/internal/store/postgres/message_send_idempotency_integration_test.go b/internal/store/postgres/message_send_idempotency_integration_test.go new file mode 100644 index 00000000..77abd4a9 --- /dev/null +++ b/internal/store/postgres/message_send_idempotency_integration_test.go @@ -0,0 +1,319 @@ +package postgres + +import ( + "context" + "errors" + "os" + "sync" + "testing" + "time" + + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgxpool" + + "telesrv/internal/domain" +) + +func TestMessageStorePrivateRandomIDConflictMatrix(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + sender := createTestUser(t, ctx, users, "+1881"+suffix+"01", "IDSender", "") + recipient := createTestUser(t, ctx, users, "+1881"+suffix+"02", "IDRecipient", "") + other := createTestUser(t, ctx, users, "+1881"+suffix+"03", "IDOther", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{sender.ID, recipient.ID, other.ID}) + }) + + messages := NewMessageStore(pool) + base := domain.SendPrivateTextRequest{ + SenderUserID: sender.ID, + RecipientUserID: recipient.ID, + RandomID: 771001, + Message: "immutable payload", + Date: 1700001000, + } + first, err := messages.SendPrivateText(ctx, base) + if err != nil { + t.Fatalf("first send: %v", err) + } + + tests := []struct { + name string + mutate func(*domain.SendPrivateTextRequest) + }{ + {name: "peer", mutate: func(req *domain.SendPrivateTextRequest) { req.RecipientUserID = other.ID }}, + {name: "body", mutate: func(req *domain.SendPrivateTextRequest) { req.Message = "different body" }}, + {name: "media", mutate: func(req *domain.SendPrivateTextRequest) { + req.Media = &domain.MessageMedia{ + Kind: domain.MessageMediaKindContact, + Contact: &domain.MessageContact{ + PhoneNumber: "+10000000000", + FirstName: "Different", + }, + } + }}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + req := base + tc.mutate(&req) + if _, err := messages.SendPrivateText(ctx, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("conflicting replay err = %v, want ErrMessageRandomIDDuplicate", err) + } + }) + } + + var privateCount, boxCount, eventCount, outboxCount int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM private_messages WHERE sender_user_id = $1 AND random_id = $2`, sender.ID, base.RandomID).Scan(&privateCount); err != nil { + t.Fatalf("count private messages: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM message_boxes WHERE private_message_id = $1`, first.SenderMessage.UID).Scan(&boxCount); err != nil { + t.Fatalf("count message boxes: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM user_update_events WHERE user_id = ANY($1::bigint[])`, []int64{sender.ID, recipient.ID, other.ID}).Scan(&eventCount); err != nil { + t.Fatalf("count update events: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM dispatch_outbox WHERE target_user_id = ANY($1::bigint[])`, []int64{sender.ID, recipient.ID, other.ID}).Scan(&outboxCount); err != nil { + t.Fatalf("count outbox: %v", err) + } + if privateCount != 1 || boxCount != 2 || eventCount != 2 || outboxCount != 2 { + t.Fatalf("rows after conflicts = private %d boxes %d events %d outbox %d, want 1/2/2/2", privateCount, boxCount, eventCount, outboxCount) + } +} + +func TestMessageStorePrivateRandomIDReplaySelfAndBlocked(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + self := createTestUser(t, ctx, users, "+1882"+suffix+"01", "IDSelf", "") + sender := createTestUser(t, ctx, users, "+1882"+suffix+"02", "BlockedSender", "") + recipient := createTestUser(t, ctx, users, "+1882"+suffix+"03", "BlockedRecipient", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{self.ID, sender.ID, recipient.ID}) + }) + + messages := NewMessageStore(pool) + selfReq := domain.SendPrivateTextRequest{ + SenderUserID: self.ID, RecipientUserID: self.ID, RandomID: 772001, + Message: "saved note", Date: 1700001100, + } + selfFirst, err := messages.SendPrivateText(ctx, selfReq) + if err != nil { + t.Fatalf("self first: %v", err) + } + selfReq.Date++ + selfReq.OriginSessionID = 99 + selfReq.RecipientBlocked = true + selfReplay, err := messages.SendPrivateText(ctx, selfReq) + if err != nil { + t.Fatalf("self replay: %v", err) + } + if !selfReplay.Duplicate || selfReplay.SenderMessage.ID != selfFirst.SenderMessage.ID || selfReplay.RecipientMessage.ID != selfFirst.SenderMessage.ID { + t.Fatalf("self replay = %+v, want original single box", selfReplay) + } + + blockedReq := domain.SendPrivateTextRequest{ + SenderUserID: sender.ID, RecipientUserID: recipient.ID, RandomID: 772002, + Message: "blocked delivery", Date: 1700001110, RecipientBlocked: true, + } + blockedFirst, err := messages.SendPrivateText(ctx, blockedReq) + if err != nil { + t.Fatalf("blocked first: %v", err) + } + if blockedFirst.RecipientMessage.ID != 0 { + t.Fatalf("blocked recipient message = %+v, want empty", blockedFirst.RecipientMessage) + } + blockedReq.Date++ + blockedReq.RecipientBlocked = false + blockedReplay, err := messages.SendPrivateText(ctx, blockedReq) + if err != nil { + t.Fatalf("blocked replay: %v", err) + } + if !blockedReplay.Duplicate || blockedReplay.SenderMessage.ID != blockedFirst.SenderMessage.ID || blockedReplay.RecipientMessage.ID != 0 { + t.Fatalf("blocked replay = %+v, want original sender-only result", blockedReplay) + } + + var selfBoxes, blockedBoxes, recipientEvents, recipientOutbox int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM message_boxes WHERE private_message_id = $1`, selfFirst.SenderMessage.UID).Scan(&selfBoxes); err != nil { + t.Fatalf("count self boxes: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM message_boxes WHERE private_message_id = $1`, blockedFirst.SenderMessage.UID).Scan(&blockedBoxes); err != nil { + t.Fatalf("count blocked boxes: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM user_update_events WHERE user_id = $1`, recipient.ID).Scan(&recipientEvents); err != nil { + t.Fatalf("count blocked recipient events: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM dispatch_outbox WHERE target_user_id = $1`, recipient.ID).Scan(&recipientOutbox); err != nil { + t.Fatalf("count blocked recipient outbox: %v", err) + } + if selfBoxes != 1 || blockedBoxes != 1 || recipientEvents != 0 || recipientOutbox != 0 { + t.Fatalf("replay rows = self boxes %d blocked boxes %d recipient events %d outbox %d, want 1/1/0/0", selfBoxes, blockedBoxes, recipientEvents, recipientOutbox) + } +} + +func TestMessageStorePrivateRandomIDReplayUsesCurrentSnapshotAndDurableDelete(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + sender := createTestUser(t, ctx, users, "+1884"+suffix+"01", "ReceiptSender", "") + recipient := createTestUser(t, ctx, users, "+1884"+suffix+"02", "ReceiptRecipient", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{sender.ID, recipient.ID}) + }) + + messages := NewMessageStore(pool) + type replayState struct { + events int + outbox int + pts int + } + loadReplayState := func() replayState { + t.Helper() + var state replayState + if err := pool.QueryRow(ctx, `SELECT count(*) FROM user_update_events WHERE user_id = $1`, sender.ID).Scan(&state.events); err != nil { + t.Fatalf("count sender events: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT count(*) FROM dispatch_outbox WHERE target_user_id = $1`, sender.ID).Scan(&state.outbox); err != nil { + t.Fatalf("count sender outbox: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT contiguous_pts FROM user_update_watermarks WHERE user_id = $1`, sender.ID).Scan(&state.pts); err != nil { + t.Fatalf("load sender pts: %v", err) + } + return state + } + req := domain.SendPrivateTextRequest{ + SenderUserID: sender.ID, RecipientUserID: recipient.ID, RandomID: 774001, + Message: "immutable receipt", Date: 1700001300, + } + first, err := messages.SendPrivateText(ctx, req) + if err != nil { + t.Fatalf("first send: %v", err) + } + edited, err := messages.EditMessage(ctx, domain.EditMessageRequest{ + OwnerUserID: sender.ID, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: recipient.ID}, + ID: first.SenderMessage.ID, + Message: "edited projection", + EditDate: 1700001301, + }) + if err != nil { + t.Fatalf("edit message: %v", err) + } + beforeReplay := loadReplayState() + replay, err := messages.SendPrivateText(ctx, req) + if err != nil { + t.Fatalf("replay after edit: %v", err) + } + if replay.SenderMessage.ID != first.SenderMessage.ID || replay.SenderMessage.Pts != edited.Self().Message.Pts || replay.SenderMessage.Body != "edited projection" || + replay.RecipientMessage.ID != first.RecipientMessage.ID || replay.RecipientMessage.Pts != first.RecipientMessage.Pts { + t.Fatalf("replay after edit = %+v/%+v, want current sender snapshot and immutable recipient receipt %d/%d", + replay.SenderMessage, replay.RecipientMessage, + first.RecipientMessage.ID, first.RecipientMessage.Pts) + } + if replay.SenderEvent.Pts != first.SenderEvent.Pts || replay.ReplayDeleteEvent != nil { + t.Fatalf("replay after edit event = %+v delete=%+v, want first-send pts and no delete", replay.SenderEvent, replay.ReplayDeleteEvent) + } + if after := loadReplayState(); after != beforeReplay { + t.Fatalf("edit replay mutated durable state = %+v, want %+v", after, beforeReplay) + } + deleted, err := messages.DeleteMessages(ctx, domain.DeleteMessagesRequest{ + OwnerUserID: sender.ID, + IDs: []int{first.SenderMessage.ID}, + Revoke: true, + Date: 1700001302, + }) + if err != nil { + t.Fatalf("delete message: %v", err) + } + beforeReplay = loadReplayState() + replay, err = messages.SendPrivateText(ctx, req) + if err != nil { + t.Fatalf("replay after delete: %v", err) + } + if replay.SenderMessage.ID != first.SenderMessage.ID || replay.SenderMessage.Pts != first.SenderMessage.Pts || replay.SenderMessage.Body != "immutable receipt" { + t.Fatalf("replay after delete = %+v, want immutable first sender snapshot", replay.SenderMessage) + } + if replay.ReplayDeleteEvent == nil || replay.ReplayDeleteEvent.Pts != deleted.Self().Event.Pts || + len(replay.ReplayDeleteEvent.MessageIDs) != 1 || replay.ReplayDeleteEvent.MessageIDs[0] != first.SenderMessage.ID { + t.Fatalf("replay delete event = %+v, want durable delete %+v", replay.ReplayDeleteEvent, deleted.Self().Event) + } + if after := loadReplayState(); after != beforeReplay { + t.Fatalf("delete replay mutated durable state = %+v, want %+v", after, beforeReplay) + } +} + +// beginHookDB lets the test commit a competing send exactly after the outer +// fast-path lookup and before its transaction starts. With MaxConns=1, a +// duplicate fallback that queries the pool while holding the transaction would +// wait until the context deadline; reading through qtx completes immediately. +type beginHookDB struct { + *pgxpool.Pool + once sync.Once + before func(context.Context) error + beforeErr error +} + +func (db *beginHookDB) Begin(ctx context.Context) (pgx.Tx, error) { + db.once.Do(func() { + if db.before != nil { + db.beforeErr = db.before(ctx) + } + }) + if db.beforeErr != nil { + return nil, db.beforeErr + } + return db.Pool.Begin(ctx) +} + +func TestMessageStorePrivateRandomIDConflictFallbackUsesTransactionConnection(t *testing.T) { + dsn := os.Getenv("TELESRV_TEST_POSTGRES_DSN") + if dsn == "" { + t.Skip("set TELESRV_TEST_POSTGRES_DSN to run postgres integration test") + } + if err := Migrate(dsn); err != nil { + t.Fatalf("migrate: %v", err) + } + config, err := pgxpool.ParseConfig(dsn) + if err != nil { + t.Fatalf("parse pool config: %v", err) + } + config.MaxConns = 1 + pool, err := pgxpool.NewWithConfig(context.Background(), config) + if err != nil { + t.Fatalf("open single-connection pool: %v", err) + } + t.Cleanup(pool.Close) + + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + sender := createTestUser(t, ctx, users, "+1883"+suffix+"01", "PoolSender", "") + recipient := createTestUser(t, ctx, users, "+1883"+suffix+"02", "PoolRecipient", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{sender.ID, recipient.ID}) + }) + + req := domain.SendPrivateTextRequest{ + SenderUserID: sender.ID, RecipientUserID: recipient.ID, RandomID: 773001, + Message: "commit between preflight and insert", Date: 1700001200, + } + boxIDs := &perUserCounterAllocator{} + db := &beginHookDB{Pool: pool} + db.before = func(ctx context.Context) error { + _, err := NewMessageStore(pool, WithMessageAllocators(boxIDs)).SendPrivateText(ctx, req) + return err + } + deadlineCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + got, err := NewMessageStore(db, WithMessageAllocators(boxIDs)).SendPrivateText(deadlineCtx, req) + if err != nil { + t.Fatalf("conflict fallback with MaxConns=1: %v", err) + } + if !got.Duplicate || got.SenderMessage.ID == 0 || got.RecipientMessage.ID == 0 { + t.Fatalf("conflict fallback result = %+v, want committed duplicate boxes", got) + } +} diff --git a/internal/store/postgres/message_send_idempotency_test.go b/internal/store/postgres/message_send_idempotency_test.go new file mode 100644 index 00000000..89dc1242 --- /dev/null +++ b/internal/store/postgres/message_send_idempotency_test.go @@ -0,0 +1,83 @@ +package postgres + +import ( + "bytes" + "testing" + + "telesrv/internal/domain" + "telesrv/internal/store" +) + +func TestPrivateSendRequestFingerprintUsesImmutableIntent(t *testing.T) { + req := domain.SendPrivateTextRequest{ + SenderUserID: 1001, + RecipientUserID: 1002, + RandomID: 99, + Message: "hello", + Silent: true, + Date: 1700000000, + OriginSessionID: 7, + RecipientBlocked: true, + } + fingerprint := func(in domain.SendPrivateTextRequest) []byte { + t.Helper() + got, err := store.PrivateSendFingerprint(in) + if err != nil { + t.Fatalf("privateSendRequestFingerprint: %v", err) + } + return got + } + + first := fingerprint(req) + replay := req + replay.Date++ + replay.OriginSessionID++ + replay.OriginAuthKeyID[0] = 9 + replay.RecipientBlocked = false + if got := fingerprint(replay); !bytes.Equal(first, got) { + t.Fatalf("execution-context-only changes altered fingerprint: %x != %x", got, first) + } + + changedPeer := req + changedPeer.RecipientUserID++ + if got := fingerprint(changedPeer); bytes.Equal(first, got) { + t.Fatal("changed recipient retained fingerprint") + } + changedBody := req + changedBody.Message = "different" + if got := fingerprint(changedBody); bytes.Equal(first, got) { + t.Fatal("changed body retained fingerprint") + } + changedMedia := req + changedMedia.Media = &domain.MessageMedia{ + Kind: domain.MessageMediaKindContact, + Contact: &domain.MessageContact{ + PhoneNumber: "+10000000000", + FirstName: "Changed", + }, + } + if got := fingerprint(changedMedia); bytes.Equal(first, got) { + t.Fatal("changed media retained fingerprint") + } +} + +func TestPrivateSendRequestFingerprintPrefersRPCFingerprint(t *testing.T) { + want := bytes.Repeat([]byte{0x5a}, 32) + req := domain.SendPrivateTextRequest{IdempotencyFingerprint: want} + got, err := store.PrivateSendFingerprint(req) + if err != nil { + t.Fatalf("privateSendRequestFingerprint: %v", err) + } + if !bytes.Equal(got, want) { + t.Fatalf("fingerprint = %x, want %x", got, want) + } + got[0] ^= 0xff + if got[0] == want[0] { + t.Fatal("returned fingerprint aliases caller storage") + } + + req.IdempotencyFingerprint = []byte{1, 2, 3} + if _, err := store.PrivateSendFingerprint(req); err == nil { + t.Fatal("short caller fingerprint accepted") + } +} diff --git a/internal/store/postgres/message_send_integration_test.go b/internal/store/postgres/message_send_integration_test.go index 4ec9cb33..6462edac 100644 --- a/internal/store/postgres/message_send_integration_test.go +++ b/internal/store/postgres/message_send_integration_test.go @@ -2,6 +2,7 @@ package postgres import ( "context" + "errors" "sync" "testing" @@ -180,15 +181,9 @@ func TestMessageStoreWebViewDataServiceActionRoundTrip(t *testing.T) { }, }, } - dup, err := messages.SendPrivateText(ctx, dupReq) - if err != nil { - t.Fatalf("SendPrivateText duplicate: %v", err) + if _, err := messages.SendPrivateText(ctx, dupReq); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("changed-media duplicate err = %v, want ErrMessageRandomIDDuplicate", err) } - if !dup.Duplicate || dup.SenderMessage.ID != got.SenderMessage.ID || dup.RecipientMessage.ID != got.RecipientMessage.ID { - t.Fatalf("duplicate = %+v, want original boxes", dup) - } - assertWebViewData("duplicate sender", dup.SenderMessage) - assertWebViewData("duplicate recipient", dup.RecipientMessage) recipientHistory, err := messages.ListByUser(ctx, recipient.ID, domain.MessageFilter{ HasPeer: true, diff --git a/internal/store/postgres/private_send_migration_compat_integration_test.go b/internal/store/postgres/private_send_migration_compat_integration_test.go new file mode 100644 index 00000000..2564e874 --- /dev/null +++ b/internal/store/postgres/private_send_migration_compat_integration_test.go @@ -0,0 +1,222 @@ +package postgres + +import ( + "context" + "errors" + "strings" + "testing" + + "github.com/jackc/pgx/v5" + + "telesrv/deploy" + "telesrv/internal/domain" + storepkg "telesrv/internal/store" +) + +func TestPrivateSendFreshMigrationsKeepLegacyWriterDefaults(t *testing.T) { + t.Parallel() + for _, test := range []struct { + migration string + adds []string + drops []string + }{ + { + migration: "migrations/0062_private_message_idempotency.up.sql", + adds: []string{ + "ADD COLUMN request_fingerprint bytea NOT NULL DEFAULT '\\x'", + "ADD COLUMN recipient_delivered boolean NOT NULL DEFAULT false", + }, + drops: []string{ + "ALTER COLUMN request_fingerprint DROP DEFAULT", + "ALTER COLUMN recipient_delivered DROP DEFAULT", + }, + }, + { + migration: "migrations/0068_private_message_send_receipt.up.sql", + adds: []string{ + "ADD COLUMN sender_box_id integer NOT NULL DEFAULT 0", + "ADD COLUMN sender_pts integer NOT NULL DEFAULT 0", + "ADD COLUMN recipient_box_id integer NOT NULL DEFAULT 0", + "ADD COLUMN recipient_pts integer NOT NULL DEFAULT 0", + }, + drops: []string{ + "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", + }, + }, + } { + t.Run(test.migration, func(t *testing.T) { + sql, err := deploy.Migrations.ReadFile(test.migration) + if err != nil { + t.Fatalf("read %s: %v", test.migration, err) + } + body := string(sql) + for _, add := range test.adds { + if !strings.Contains(body, add) { + t.Errorf("%s does not install legacy-writer default %q", test.migration, add) + } + } + for _, drop := range test.drops { + if strings.Contains(body, drop) { + t.Errorf("%s removes legacy-writer default %q", test.migration, drop) + } + } + }) + } +} + +func TestPrivateSendMigration75To77PreservesOldWritersAndRejectsUnknownReceiptsPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + sender := createTestUser(t, ctx, users, "+1888"+suffix+"01", "MigrationSender", "") + recipient := createTestUser(t, ctx, users, "+1888"+suffix+"02", "MigrationRecipient", "") + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{sender.ID, recipient.ID}) + }) + + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin migration compatibility tx: %v", err) + } + defer func() { _ = tx.Rollback(context.Background()) }() + + // Recreate the schema shape left by the original 0062/0068 migrations at + // version 75, then apply the corrective expand migration in isolation. + if _, err := tx.Exec(ctx, ` +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`); err != nil { + t.Fatalf("simulate version 75 defaults: %v", err) + } + upSQL, err := deploy.Migrations.ReadFile("migrations/0077_correct_private_send_defaults.up.sql") + if err != nil { + t.Fatalf("read 0077 up: %v", err) + } + if _, err := tx.Exec(ctx, string(upSQL)); err != nil { + t.Fatalf("apply 0077 up: %v", err) + } + + var defaultCount int + if err := tx.QueryRow(ctx, ` +SELECT count(*) +FROM information_schema.columns +WHERE table_schema = 'public' + AND table_name = 'private_messages' + AND column_name = ANY($1::text[]) + AND column_default IS NOT NULL`, []string{ + "request_fingerprint", "recipient_delivered", "sender_box_id", + "sender_pts", "recipient_box_id", "recipient_pts", + }).Scan(&defaultCount); err != nil { + t.Fatalf("inspect restored defaults: %v", err) + } + if defaultCount != 6 { + t.Fatalf("restored private-send defaults = %d, want 6", defaultCount) + } + + insertLegacyBoxes := func(privateMessageID int64, boxID, pts, date int, body string) { + t.Helper() + if _, err := tx.Exec(ctx, ` +INSERT INTO message_boxes ( + owner_user_id, box_id, private_message_id, message_sender_id, + peer_type, peer_id, from_user_id, message_date, outgoing, body, entities, pts +) VALUES + ($1, $3, $5, $1, 'user', $2, $1, $6, true, $7, '[]'::jsonb, $4), + ($2, $3, $5, $1, 'user', $1, $1, $6, false, $7, '[]'::jsonb, $4)`, + sender.ID, recipient.ID, boxID, pts, privateMessageID, date, body); err != nil { + t.Fatalf("insert legacy message boxes: %v", err) + } + } + + pre0062 := domain.SendPrivateTextRequest{ + SenderUserID: sender.ID, RecipientUserID: recipient.ID, + RandomID: 887_001, Message: "pre-0062 writer", Date: 1_700_040_001, + } + var pre0062ID int64 + if err := tx.QueryRow(ctx, ` +INSERT INTO private_messages ( + sender_user_id, recipient_user_id, random_id, message_date, body, entities +) VALUES ($1, $2, $3, $4, $5, '[]'::jsonb) +RETURNING id`, pre0062.SenderUserID, pre0062.RecipientUserID, pre0062.RandomID, pre0062.Date, pre0062.Message).Scan(&pre0062ID); err != nil { + t.Fatalf("pre-0062 INSERT after migration: %v", err) + } + insertLegacyBoxes(pre0062ID, 1, 1, pre0062.Date, pre0062.Message) + assertLegacyPrivateSendSentinels(t, ctx, tx, pre0062ID, false) + if _, err := NewMessageStore(tx).SendPrivateText(ctx, pre0062); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) { + t.Fatalf("pre-0062 unknown replay err = %v, want ErrMessageRandomIDDuplicate", err) + } + + era0062 := domain.SendPrivateTextRequest{ + SenderUserID: sender.ID, RecipientUserID: recipient.ID, + RandomID: 887_002, Message: "0062-era writer", Date: 1_700_040_002, + } + fingerprint, err := storepkg.PrivateSendFingerprint(era0062) + if err != nil { + t.Fatalf("fingerprint 0062-era request: %v", err) + } + var era0062ID int64 + if err := tx.QueryRow(ctx, ` +INSERT INTO private_messages ( + sender_user_id, recipient_user_id, random_id, request_fingerprint, + recipient_delivered, message_date, body, entities +) VALUES ($1, $2, $3, $4, true, $5, $6, '[]'::jsonb) +RETURNING id`, era0062.SenderUserID, era0062.RecipientUserID, era0062.RandomID, + fingerprint, era0062.Date, era0062.Message).Scan(&era0062ID); err != nil { + t.Fatalf("0062-era INSERT after migration: %v", err) + } + insertLegacyBoxes(era0062ID, 2, 2, era0062.Date, era0062.Message) + assertLegacyPrivateSendSentinels(t, ctx, tx, era0062ID, true) + if _, err := NewMessageStore(tx).SendPrivateText(ctx, era0062); err == nil || + errors.Is(err, domain.ErrMessageRandomIDDuplicate) || + !strings.Contains(err.Error(), "invalid immutable sender receipt") { + t.Fatalf("0062-era unknown receipt err = %v, want explicit invalid receipt failure", err) + } + + var privateRows, eventRows, outboxRows int + if err := tx.QueryRow(ctx, `SELECT count(*) FROM private_messages WHERE sender_user_id = $1`, sender.ID).Scan(&privateRows); err != nil { + t.Fatalf("count legacy private rows: %v", err) + } + if err := tx.QueryRow(ctx, `SELECT count(*) FROM user_update_events WHERE user_id = ANY($1::bigint[])`, []int64{sender.ID, recipient.ID}).Scan(&eventRows); err != nil { + t.Fatalf("count legacy replay events: %v", err) + } + if err := tx.QueryRow(ctx, `SELECT count(*) FROM dispatch_outbox WHERE target_user_id = ANY($1::bigint[])`, []int64{sender.ID, recipient.ID}).Scan(&outboxRows); err != nil { + t.Fatalf("count legacy replay outbox: %v", err) + } + if privateRows != 2 || eventRows != 0 || outboxRows != 0 { + t.Fatalf("legacy replay facts private/events/outbox = %d/%d/%d, want 2/0/0", privateRows, eventRows, outboxRows) + } +} + +func assertLegacyPrivateSendSentinels(t *testing.T, ctx context.Context, q pgx.Tx, id int64, delivered bool) { + t.Helper() + var ( + fingerprint []byte + gotDelivered bool + senderBoxID, senderPts, recipientBoxID, recipientPts int + ) + if err := q.QueryRow(ctx, ` +SELECT request_fingerprint, recipient_delivered, + sender_box_id, sender_pts, recipient_box_id, recipient_pts +FROM private_messages +WHERE id = $1`, id).Scan( + &fingerprint, &gotDelivered, + &senderBoxID, &senderPts, &recipientBoxID, &recipientPts, + ); err != nil { + t.Fatalf("load legacy private-send sentinels: %v", err) + } + if (!delivered && (len(fingerprint) != 0 || gotDelivered)) || + senderBoxID != 0 || senderPts != 0 || recipientBoxID != 0 || recipientPts != 0 { + t.Fatalf("legacy sentinels fingerprint=%x delivered=%v receipt=%d/%d/%d/%d", + fingerprint, gotDelivered, senderBoxID, senderPts, recipientBoxID, recipientPts) + } + if delivered && (len(fingerprint) != 32 || !gotDelivered) { + t.Fatalf("0062-era fingerprint/delivery = %x/%v, want 32 bytes/true", fingerprint, gotDelivered) + } +} diff --git a/internal/store/postgres/queries/message.sql b/internal/store/postgres/queries/message.sql index dc22f777..bfc8e284 100644 --- a/internal/store/postgres/queries/message.sql +++ b/internal/store/postgres/queries/message.sql @@ -4,6 +4,12 @@ WITH pm AS ( sender_user_id, recipient_user_id, random_id, + request_fingerprint, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts, message_date, body, entities @@ -11,6 +17,12 @@ WITH pm AS ( sqlc.arg(from_user_id), sqlc.arg(owner_user_id), 0, + '\x'::bytea, + false, + 0, + 0, + 0, + 0, sqlc.arg(message_date), sqlc.arg(body), sqlc.arg(entities_json)::jsonb @@ -82,6 +94,12 @@ INSERT INTO private_messages ( sender_user_id, recipient_user_id, random_id, + request_fingerprint, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts, message_date, ttl_period, expires_at, @@ -108,7 +126,9 @@ INSERT INTO private_messages ( grouped_id, effect ) VALUES ( - $1, $2, $3, $4, sqlc.arg(ttl_period)::int, sqlc.arg(expires_at)::int, $5, sqlc.arg(entities_json)::jsonb, + $1, $2, $3, sqlc.arg(request_fingerprint)::bytea, sqlc.arg(recipient_delivered)::boolean, + 0, 0, 0, 0, + $4, sqlc.arg(ttl_period)::int, sqlc.arg(expires_at)::int, $5, sqlc.arg(entities_json)::jsonb, sqlc.arg(silent)::boolean, sqlc.arg(noforwards)::boolean, sqlc.arg(reply_to_msg_id)::int, @@ -149,6 +169,17 @@ SELECT sender_user_id, recipient_user_id, random_id, + request_fingerprint, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts, + sender_snapshot::text AS sender_snapshot_json, + sender_delete_pts, + sender_delete_pts_count, + sender_delete_date, + sender_delete_message_ids::text AS sender_delete_message_ids_json, message_date, ttl_period, expires_at, diff --git a/internal/store/postgres/queries/user_update_event.sql b/internal/store/postgres/queries/user_update_event.sql index b5b4bb5d..230b49e8 100644 --- a/internal/store/postgres/queries/user_update_event.sql +++ b/internal/store/postgres/queries/user_update_event.sql @@ -208,29 +208,31 @@ INSERT INTO dispatch_outbox ( ON CONFLICT DO NOTHING; -- name: ClaimDispatchOutbox :many -WITH picked AS ( - SELECT d.target_user_id, d.id - FROM dispatch_outbox d +-- durable head 表只保留每用户一行，并同步 head 的 readiness。claim 先锁 +-- lane head 再更新对应 outbox 行，既不会扫描 backlog，也不会并发领取同一用户。 +WITH picked_heads AS ( + SELECT h.target_user_id, h.head_id + FROM dispatch_outbox_user_heads h WHERE ( - d.status = 'pending' - AND d.next_attempt_at <= now() - ) - OR ( - d.status = 'dispatching' - AND d.updated_at < now() - make_interval(secs => sqlc.arg(lease_seconds)::int) - ) - ORDER BY d.next_attempt_at ASC, d.target_user_id ASC, d.pts ASC, d.id ASC + h.status = 'pending' + AND h.next_attempt_at <= now() + ) + OR ( + h.status = 'dispatching' + AND h.updated_at < now() - make_interval(secs => sqlc.arg(lease_seconds)::int) + ) + ORDER BY h.next_attempt_at ASC, h.target_user_id ASC, h.head_pts ASC, h.head_id ASC LIMIT sqlc.arg(limit_count) - FOR UPDATE SKIP LOCKED + FOR UPDATE OF h SKIP LOCKED ) UPDATE dispatch_outbox d SET status = 'dispatching', attempts = d.attempts + 1, updated_at = now() -FROM picked p +FROM picked_heads p WHERE d.target_user_id = p.target_user_id - AND d.id = p.id + AND d.id = p.head_id RETURNING d.id, d.target_user_id, @@ -240,25 +242,85 @@ RETURNING d.exclude_session_id, d.attempts; --- name: MarkDispatchDelivered :exec +-- name: ClaimDispatchOutboxShards :many +-- 固定 logical shard 由 target_user_id 决定；运行时 worker 只领取分配给自己的 +-- shard 集合，因此同一用户永远只有一条串行 lane，而不同用户可并行。 +WITH picked_heads AS ( + SELECT h.target_user_id, h.head_id + FROM dispatch_outbox_user_heads h + -- 256 与 store.DispatchOutboxLogicalShards、0069 generated column 是同一 + -- schema 常量；不得随 worker 数变化。 + WHERE h.logical_shard = ANY(sqlc.arg(shard_ids)::smallint[]) + AND ( + ( + h.status = 'pending' + AND h.next_attempt_at <= now() + ) + OR ( + h.status = 'dispatching' + AND h.updated_at < now() - make_interval(secs => sqlc.arg(lease_seconds)::int) + ) + ) + ORDER BY h.next_attempt_at ASC, h.target_user_id ASC, h.head_pts ASC, h.head_id ASC + LIMIT sqlc.arg(limit_count) + FOR UPDATE OF h SKIP LOCKED +) +UPDATE dispatch_outbox d +SET + status = 'dispatching', + attempts = d.attempts + 1, + updated_at = now() +FROM picked_heads p +WHERE d.target_user_id = p.target_user_id + AND d.id = p.head_id +RETURNING + d.id, + d.target_user_id, + d.pts, + d.event_type, + d.exclude_auth_key_id, + d.exclude_session_id, + d.attempts; + +-- name: MarkDispatchDelivered :execrows -- 方案 A：投递成功即删除。outbox 是任务队列，delivered 行无保留价值 -- （消息在 message_boxes、离线补偿在 user_update_events），删除让表维持「未完成任务」小稳态。 -DELETE FROM dispatch_outbox -WHERE target_user_id = $1 - AND id = $2; +-- claim 的锁序是 user_heads→outbox；completion 必须先显式锁同一 head 再删 outbox， +-- 否则租约过期 claim 与完成恰好竞争时会形成 outbox→head / head→outbox 环路。 +WITH locked_head AS MATERIALIZED ( + SELECT h.target_user_id + FROM dispatch_outbox_user_heads h + WHERE h.target_user_id = sqlc.arg(target_user_id)::bigint + FOR UPDATE +) +DELETE FROM dispatch_outbox d +USING locked_head h +WHERE d.target_user_id = h.target_user_id + AND d.id = sqlc.arg(id)::bigint + AND d.status = 'dispatching' + AND d.attempts = sqlc.arg(expected_attempts)::int; --- name: MarkDispatchFailed :exec -UPDATE dispatch_outbox +-- name: MarkDispatchFailed :execrows +WITH locked_head AS MATERIALIZED ( + SELECT h.target_user_id + FROM dispatch_outbox_user_heads h + WHERE h.target_user_id = sqlc.arg(target_user_id)::bigint + FOR UPDATE +) +UPDATE dispatch_outbox d SET - status = CASE WHEN attempts >= 5 THEN 'failed' ELSE 'pending' END, + status = CASE WHEN d.attempts >= 5 THEN 'failed' ELSE 'pending' END, next_attempt_at = CASE - WHEN attempts >= 5 THEN next_attempt_at - ELSE now() + make_interval(secs => LEAST(60, attempts * attempts)) + WHEN d.attempts >= 5 THEN d.next_attempt_at + ELSE now() + make_interval(secs => LEAST(60, d.attempts * d.attempts)) END, - last_error = $3, + last_error = sqlc.arg(last_error)::text, updated_at = now() -WHERE target_user_id = $1 - AND id = $2; +FROM locked_head h +WHERE d.target_user_id = h.target_user_id + AND d.id = sqlc.arg(id)::bigint + AND d.status = 'dispatching' + AND d.attempts = sqlc.arg(expected_attempts)::int; -- name: BatchListDispatchEvents :many -- 按 (user_id, pts) 精确批量取账号事件，供 outbox worker 一次性加载一批 claim 的事件详情， @@ -394,22 +456,44 @@ LEFT JOIN users from_u ON from_u.id = m.from_user_id LEFT JOIN users fwd_u ON m.fwd_from_peer_type = 'user' AND fwd_u.id = m.fwd_from_peer_id LEFT JOIN users reply_u ON m.reply_to_peer_type = 'user' AND reply_u.id = m.reply_to_peer_id; --- name: MarkDispatchDeliveredBatch :exec +-- name: MarkDispatchDeliveredBatch :execrows -- 批量删除一批已投递的 (target_user_id, id)；target_user_id 入 WHERE 命中唯一索引并避免串删。 +WITH input AS MATERIALIZED ( + SELECT tu.target_user_id, di.id, ea.attempts + FROM unnest(@target_user_ids::bigint[]) WITH ORDINALITY AS tu(target_user_id, ord) + JOIN unnest(@ids::bigint[]) WITH ORDINALITY AS di(id, ord) USING (ord) + JOIN unnest(@expected_attempts::int[]) WITH ORDINALITY AS ea(attempts, ord) USING (ord) +), +locked_heads AS MATERIALIZED ( + SELECT h.target_user_id + FROM dispatch_outbox_user_heads h + JOIN (SELECT DISTINCT target_user_id FROM input) i USING (target_user_id) + -- Match ClaimDispatchOutbox[Shards] exactly. A stale-lease claim may lock several + -- dispatching heads while this completion batch locks the same set; a different + -- multi-row order would merely move the deadlock one level up. + ORDER BY h.next_attempt_at, h.target_user_id, h.head_pts, h.head_id + FOR UPDATE OF h +) DELETE FROM dispatch_outbox d -USING unnest(@target_user_ids::bigint[]) WITH ORDINALITY AS tu(target_user_id, ord) -JOIN unnest(@ids::bigint[]) WITH ORDINALITY AS di(id, ord) USING (ord) -WHERE d.target_user_id = tu.target_user_id - AND d.id = di.id; +USING input i, locked_heads h +WHERE d.target_user_id = h.target_user_id + AND d.target_user_id = i.target_user_id + AND d.id = i.id + AND d.status = 'dispatching' + AND d.attempts = i.attempts; -- name: DeleteFailedDispatchOutbox :one -WITH doomed AS ( - SELECT target_user_id, id - FROM dispatch_outbox - WHERE status = 'failed' - AND updated_at < now() - make_interval(secs => sqlc.arg(older_than_seconds)::int) - ORDER BY updated_at ASC, target_user_id ASC, id ASC +-- failed 只能成为 lane head；从 head 表开始并先锁 head，既走 0074 的小索引，也与 +-- claim/completion 保持同一 user_heads→outbox 锁序。删除的只是在线任务，durable +-- user_update_events 不动，故客户端仍可经 difference 恢复。 +WITH doomed AS MATERIALIZED ( + SELECT h.target_user_id, h.head_id AS id + FROM dispatch_outbox_user_heads h + WHERE h.status = 'failed' + AND h.updated_at < now() - make_interval(secs => sqlc.arg(older_than_seconds)::int) + ORDER BY h.updated_at ASC, h.target_user_id ASC, h.head_id ASC LIMIT sqlc.arg(limit_count) + FOR UPDATE OF h SKIP LOCKED ), deleted AS ( DELETE FROM dispatch_outbox d diff --git a/internal/store/postgres/retention_migration_integration_test.go b/internal/store/postgres/retention_migration_integration_test.go new file mode 100644 index 00000000..f108f39f --- /dev/null +++ b/internal/store/postgres/retention_migration_integration_test.go @@ -0,0 +1,137 @@ +package postgres + +import ( + "context" + "errors" + "testing" + + "github.com/jackc/pgx/v5/pgconn" + + "telesrv/deploy" + "telesrv/internal/domain" +) + +func TestRetentionDownMigrationsRejectAdvancedFloorsPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + userID := createRevokeTestUser(t, ctx, pool, "retention-down-guard") + channel, err := NewChannelStore(pool).CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: userID, + Title: "retention down guard", + Megagroup: true, + Date: 1_700_030_000, + }) + if err != nil { + t.Fatalf("create guarded channel: %v", err) + } + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM channels WHERE id = $1", channel.Channel.ID) + }) + + for _, test := range []struct { + name string + migration string + advance func(context.Context, interface { + Exec(context.Context, string, ...any) (pgconn.CommandTag, error) + }) error + }{ + { + name: "channel", + migration: "migrations/0063_channel_update_retention.down.sql", + advance: func(ctx context.Context, tx interface { + Exec(context.Context, string, ...any) (pgconn.CommandTag, error) + }) error { + _, err := tx.Exec(ctx, ` +UPDATE channel_update_checkpoints +SET retained_through_pts = 1 +WHERE channel_id = $1`, channel.Channel.ID) + return err + }, + }, + { + name: "user", + migration: "migrations/0064_user_update_retention.down.sql", + advance: func(ctx context.Context, tx interface { + Exec(context.Context, string, ...any) (pgconn.CommandTag, error) + }) error { + _, err := tx.Exec(ctx, ` +INSERT INTO user_update_retention (user_id, retained_through_pts, retained_through_date) +VALUES ($1, 1, 1) +ON CONFLICT (user_id) DO UPDATE SET + retained_through_pts = 1, + retained_through_date = 1`, userID) + return err + }, + }, + } { + t.Run(test.name, func(t *testing.T) { + downSQL, err := deploy.Migrations.ReadFile(test.migration) + if err != nil { + t.Fatalf("read %s: %v", test.migration, err) + } + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin guarded down migration: %v", err) + } + defer func() { _ = tx.Rollback(context.Background()) }() + if err := test.advance(ctx, tx); err != nil { + t.Fatalf("advance retained floor: %v", err) + } + if _, err := tx.Exec(ctx, string(downSQL)); err == nil { + t.Fatalf("%s succeeded with retained floor > 0", test.migration) + } else { + var pgErr *pgconn.PgError + if !errors.As(err, &pgErr) || pgErr.Code != "55000" { + t.Fatalf("%s error = %v, want SQLSTATE 55000", test.migration, err) + } + } + }) + } +} + +func TestPerformanceMigrationIndexesPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + + // 0072 must reuse the base schema's unique key for its durable-head FK. A second identical + // unique index would add outbox enqueue/delete write amplification without improving lookup. + // Likewise 0067 supersedes the transitional created_at orphan-GC index with last_used_at. + var ( + baseOutboxUnique, duplicateOutboxUnique bool + lastUsedAuthIndex, obsoleteAuthIndex bool + pendingHeadIndex, staleHeadIndex bool + poisonHeadIndex bool + tempExpiryIndex bool + ) + if err := pool.QueryRow(ctx, ` +SELECT + to_regclass('public.dispatch_outbox_target_user_id_id_key') IS NOT NULL, + to_regclass('public.dispatch_outbox_target_id_uidx') IS NOT NULL, + to_regclass('public.auth_keys_orphan_last_used_idx') IS NOT NULL, + to_regclass('public.auth_keys_orphan_retention_idx') IS NOT NULL, + to_regclass('public.dispatch_outbox_user_heads_pending_shard_idx') IS NOT NULL, + to_regclass('public.dispatch_outbox_user_heads_dispatching_shard_idx') IS NOT NULL, + to_regclass('public.dispatch_outbox_user_heads_failed_cleanup_idx') IS NOT NULL, + to_regclass('public.temp_auth_key_bindings_expiry_idx') IS NOT NULL +`).Scan( + &baseOutboxUnique, + &duplicateOutboxUnique, + &lastUsedAuthIndex, + &obsoleteAuthIndex, + &pendingHeadIndex, + &staleHeadIndex, + &poisonHeadIndex, + &tempExpiryIndex, + ); err != nil { + t.Fatalf("inspect performance migration indexes: %v", err) + } + if !baseOutboxUnique || duplicateOutboxUnique { + t.Fatalf("outbox target/id indexes base=%v duplicate=%v, want true/false", baseOutboxUnique, duplicateOutboxUnique) + } + if !lastUsedAuthIndex || obsoleteAuthIndex { + t.Fatalf("auth orphan indexes last_used=%v created_at=%v, want true/false", lastUsedAuthIndex, obsoleteAuthIndex) + } + if !pendingHeadIndex || !staleHeadIndex || !poisonHeadIndex || !tempExpiryIndex { + t.Fatalf("ready/expiry indexes pending=%v stale=%v poison=%v temp_expiry=%v, want all true", pendingHeadIndex, staleHeadIndex, poisonHeadIndex, tempExpiryIndex) + } +} diff --git a/internal/store/postgres/saved_dialog_integration_test.go b/internal/store/postgres/saved_dialog_integration_test.go index 122ba099..7239b905 100644 --- a/internal/store/postgres/saved_dialog_integration_test.go +++ b/internal/store/postgres/saved_dialog_integration_test.go @@ -207,9 +207,17 @@ func TestSavedDialogsBackfillRule(t *testing.T) { t.Helper() var privateID int64 if err := pool.QueryRow(ctx, ` -INSERT INTO private_messages (sender_user_id, recipient_user_id, random_id, message_date, body, entities) -VALUES ($1, $1, $2::bigint, 1700000800, 'legacy', '[]'::jsonb) -RETURNING id`, owner.ID, 841100+boxID).Scan(&privateID); err != nil { +INSERT INTO private_messages ( + sender_user_id, recipient_user_id, random_id, request_fingerprint, recipient_delivered, + sender_box_id, sender_pts, recipient_box_id, recipient_pts, + message_date, body, entities +) +VALUES ( + $1, $1, $2::bigint, decode(repeat('00', 32), 'hex'), false, + $3::int, $3::int, 0, 0, + 1700000800, 'legacy', '[]'::jsonb +) +RETURNING id`, owner.ID, 841100+boxID, boxID).Scan(&privateID); err != nil { t.Fatalf("insert legacy private message %d: %v", boxID, err) } if _, err := pool.Exec(ctx, ` diff --git a/internal/store/postgres/send_replay_migration_integration_test.go b/internal/store/postgres/send_replay_migration_integration_test.go new file mode 100644 index 00000000..69a77787 --- /dev/null +++ b/internal/store/postgres/send_replay_migration_integration_test.go @@ -0,0 +1,45 @@ +package postgres + +import ( + "context" + "testing" + + "telesrv/deploy" +) + +func TestSendReplaySnapshotMigrationRoundTripPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + downSQL, err := deploy.Migrations.ReadFile("migrations/0073_send_replay_snapshots.down.sql") + if err != nil { + t.Fatalf("read 0073 down: %v", err) + } + upSQL, err := deploy.Migrations.ReadFile("migrations/0073_send_replay_snapshots.up.sql") + if err != nil { + t.Fatalf("read 0073 up: %v", err) + } + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin migration round trip: %v", err) + } + defer func() { _ = tx.Rollback(context.Background()) }() + if _, err := tx.Exec(ctx, string(downSQL)); err != nil { + t.Fatalf("0073 down: %v", err) + } + if _, err := tx.Exec(ctx, string(upSQL)); err != nil { + t.Fatalf("0073 up: %v", err) + } + var privateSnapshot, channelSnapshot, privateDeleteIDs, channelDeleteIDs bool + if err := tx.QueryRow(ctx, ` +SELECT + EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema='public' AND table_name='private_messages' AND column_name='sender_snapshot'), + EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema='public' AND table_name='channel_messages' AND column_name='send_snapshot'), + EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema='public' AND table_name='private_messages' AND column_name='sender_delete_message_ids'), + EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema='public' AND table_name='channel_messages' AND column_name='delete_message_ids') +`).Scan(&privateSnapshot, &channelSnapshot, &privateDeleteIDs, &channelDeleteIDs); err != nil { + t.Fatalf("inspect 0073 columns: %v", err) + } + if !privateSnapshot || !channelSnapshot || !privateDeleteIDs || !channelDeleteIDs { + t.Fatalf("0073 columns private/channel snapshot=%v/%v delete_ids=%v/%v, want all true", privateSnapshot, channelSnapshot, privateDeleteIDs, channelDeleteIDs) + } +} diff --git a/internal/store/postgres/sqlcgen/message.sql.go b/internal/store/postgres/sqlcgen/message.sql.go index abae0092..0c22293c 100644 --- a/internal/store/postgres/sqlcgen/message.sql.go +++ b/internal/store/postgres/sqlcgen/message.sql.go @@ -84,6 +84,12 @@ WITH pm AS ( sender_user_id, recipient_user_id, random_id, + request_fingerprint, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts, message_date, body, entities @@ -91,6 +97,12 @@ WITH pm AS ( $1, $2, 0, + '\x'::bytea, + false, + 0, + 0, + 0, + 0, $3, $4, $5::jsonb @@ -530,6 +542,12 @@ INSERT INTO private_messages ( sender_user_id, recipient_user_id, random_id, + request_fingerprint, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts, message_date, ttl_period, expires_at, @@ -556,27 +574,29 @@ INSERT INTO private_messages ( grouped_id, effect ) VALUES ( - $1, $2, $3, $4, $6::int, $7::int, $5, $8::jsonb, - $9::boolean, - $10::boolean, - $11::int, - $12::text, - $13::bigint, - $14::int, - $15::int, - $16::text, - $17::jsonb, - $18::int, - $19::text, - $20::bigint, + $1, $2, $3, $6::bytea, $7::boolean, + 0, 0, 0, 0, + $4, $8::int, $9::int, $5, $10::jsonb, + $11::boolean, + $12::boolean, + $13::int, + $14::text, + $15::bigint, + $16::int, + $17::int, + $18::text, + $19::jsonb, + $20::int, $21::text, - $22::int, - $23::jsonb, - $24::jsonb, + $22::bigint, + $23::text, + $24::int, $25::jsonb, - $26::bigint, - $27::bigint, - $28::bigint + $26::jsonb, + $27::jsonb, + $28::bigint, + $29::bigint, + $30::bigint ) ON CONFLICT (sender_user_id, random_id) WHERE random_id <> 0 DO NOTHING RETURNING @@ -593,34 +613,36 @@ RETURNING ` type CreatePrivateMessageParams struct { - SenderUserID int64 - RecipientUserID int64 - RandomID int64 - MessageDate int32 - Body string - TtlPeriod int32 - ExpiresAt int32 - EntitiesJson []byte - Silent bool - Noforwards bool - ReplyToMsgID int32 - ReplyToPeerType string - ReplyToPeerID int64 - ReplyToTopID int32 - ReplyToStoryID int32 - QuoteText string - QuoteEntitiesJson []byte - QuoteOffset int32 - FwdFromPeerType string - FwdFromPeerID int64 - FwdFromName string - FwdDate int32 - MediaJson []byte - ReplyMarkupJson []byte - RichMessageJson []byte - ViaBotID int64 - GroupedID int64 - Effect int64 + SenderUserID int64 + RecipientUserID int64 + RandomID int64 + MessageDate int32 + Body string + RequestFingerprint []byte + RecipientDelivered bool + TtlPeriod int32 + ExpiresAt int32 + EntitiesJson []byte + Silent bool + Noforwards bool + ReplyToMsgID int32 + ReplyToPeerType string + ReplyToPeerID int64 + ReplyToTopID int32 + ReplyToStoryID int32 + QuoteText string + QuoteEntitiesJson []byte + QuoteOffset int32 + FwdFromPeerType string + FwdFromPeerID int64 + FwdFromName string + FwdDate int32 + MediaJson []byte + ReplyMarkupJson []byte + RichMessageJson []byte + ViaBotID int64 + GroupedID int64 + Effect int64 } type CreatePrivateMessageRow struct { @@ -643,6 +665,8 @@ func (q *Queries) CreatePrivateMessage(ctx context.Context, arg CreatePrivateMes arg.RandomID, arg.MessageDate, arg.Body, + arg.RequestFingerprint, + arg.RecipientDelivered, arg.TtlPeriod, arg.ExpiresAt, arg.EntitiesJson, @@ -1969,6 +1993,17 @@ SELECT sender_user_id, recipient_user_id, random_id, + request_fingerprint, + recipient_delivered, + sender_box_id, + sender_pts, + recipient_box_id, + recipient_pts, + sender_snapshot::text AS sender_snapshot_json, + sender_delete_pts, + sender_delete_pts_count, + sender_delete_date, + sender_delete_message_ids::text AS sender_delete_message_ids_json, message_date, ttl_period, expires_at, @@ -1987,16 +2022,27 @@ type GetPrivateMessageByRandomIDParams struct { } type GetPrivateMessageByRandomIDRow struct { - ID int64 - SenderUserID int64 - RecipientUserID int64 - RandomID int64 - MessageDate int32 - TtlPeriod int32 - ExpiresAt int32 - EditDate int32 - Body string - EntitiesJson string + ID int64 + SenderUserID int64 + RecipientUserID int64 + RandomID int64 + RequestFingerprint []byte + RecipientDelivered bool + SenderBoxID int32 + SenderPts int32 + RecipientBoxID int32 + RecipientPts int32 + SenderSnapshotJson string + SenderDeletePts int32 + SenderDeletePtsCount int32 + SenderDeleteDate int32 + SenderDeleteMessageIdsJson string + MessageDate int32 + TtlPeriod int32 + ExpiresAt int32 + EditDate int32 + Body string + EntitiesJson string } func (q *Queries) GetPrivateMessageByRandomID(ctx context.Context, arg GetPrivateMessageByRandomIDParams) (GetPrivateMessageByRandomIDRow, error) { @@ -2007,6 +2053,17 @@ func (q *Queries) GetPrivateMessageByRandomID(ctx context.Context, arg GetPrivat &i.SenderUserID, &i.RecipientUserID, &i.RandomID, + &i.RequestFingerprint, + &i.RecipientDelivered, + &i.SenderBoxID, + &i.SenderPts, + &i.RecipientBoxID, + &i.RecipientPts, + &i.SenderSnapshotJson, + &i.SenderDeletePts, + &i.SenderDeletePtsCount, + &i.SenderDeleteDate, + &i.SenderDeleteMessageIdsJson, &i.MessageDate, &i.TtlPeriod, &i.ExpiresAt, diff --git a/internal/store/postgres/sqlcgen/models.go b/internal/store/postgres/sqlcgen/models.go index d8906579..71e73d9d 100644 --- a/internal/store/postgres/sqlcgen/models.go +++ b/internal/store/postgres/sqlcgen/models.go @@ -174,6 +174,7 @@ type AuthKey struct { SystemVersion string ApiID int32 AppVersion string + LastUsedAt pgtype.Timestamptz } type Authorization struct { @@ -621,6 +622,11 @@ type ChannelMessage struct { GroupedID int64 SavedPeerType string SavedPeerID int64 + SendSnapshot []byte + DeletePts int32 + DeletePtsCount int32 + DeleteDate int32 + DeleteMessageIds []byte } type ChannelMessageMedium struct { @@ -689,6 +695,14 @@ type ChannelUnreadMentionIndex struct { CreatedAt pgtype.Timestamptz } +type ChannelUpdateCheckpoint struct { + ChannelID int64 + RetainedThroughPts int32 + LatestEventDate int32 + LatestPts int32 + UpdatedAt pgtype.Timestamptz +} + type ChannelUpdateEvent struct { ChannelID int64 Pts int32 @@ -835,6 +849,16 @@ type DispatchOutbox struct { UpdatedAt pgtype.Timestamptz } +type DispatchOutboxUserHead struct { + TargetUserID int64 + HeadID int64 + HeadPts int32 + LogicalShard *int16 + Status string + NextAttemptAt pgtype.Timestamptz + UpdatedAt pgtype.Timestamptz +} + type Document struct { ID int64 AccessHash int64 @@ -1160,38 +1184,49 @@ type PrivateMediaCategoryCount struct { } type PrivateMessage struct { - ID int64 - SenderUserID int64 - RecipientUserID int64 - RandomID int64 - MessageDate int32 - Body string - Entities []byte - CreatedAt pgtype.Timestamptz - EditDate int32 - Silent bool - Noforwards bool - ReplyToMsgID int32 - ReplyToPeerType string - ReplyToPeerID int64 - ReplyToTopID int32 - QuoteText string - QuoteEntities []byte - QuoteOffset int32 - FwdFromPeerType string - FwdFromPeerID int64 - FwdFromName string - FwdDate int32 - Media []byte - TtlPeriod int32 - ExpiresAt int32 - ReplyMarkup []byte - ViaBotID int64 - RichMessage []byte - GroupedID int64 - ReplyToStoryID int32 - Effect int64 - HideEdited bool + ID int64 + SenderUserID int64 + RecipientUserID int64 + RandomID int64 + MessageDate int32 + Body string + Entities []byte + CreatedAt pgtype.Timestamptz + EditDate int32 + Silent bool + Noforwards bool + ReplyToMsgID int32 + ReplyToPeerType string + ReplyToPeerID int64 + ReplyToTopID int32 + QuoteText string + QuoteEntities []byte + QuoteOffset int32 + FwdFromPeerType string + FwdFromPeerID int64 + FwdFromName string + FwdDate int32 + Media []byte + TtlPeriod int32 + ExpiresAt int32 + ReplyMarkup []byte + ViaBotID int64 + RichMessage []byte + GroupedID int64 + ReplyToStoryID int32 + Effect int64 + HideEdited bool + RequestFingerprint []byte + RecipientDelivered bool + SenderBoxID int32 + SenderPts int32 + RecipientBoxID int32 + RecipientPts int32 + SenderSnapshot []byte + SenderDeletePts int32 + SenderDeletePtsCount int32 + SenderDeleteDate int32 + SenderDeleteMessageIds []byte } type PrivateMessageReaction struct { @@ -1482,13 +1517,14 @@ type ThemeUserInstall struct { } type UpdateState struct { - AuthKeyID int64 - Pts int32 - Qts int32 - Date int32 - Seq int32 - UpdatedAt pgtype.Timestamptz - UserID int64 + AuthKeyID int64 + Pts int32 + Qts int32 + Date int32 + Seq int32 + UpdatedAt pgtype.Timestamptz + UserID int64 + ObservedPts int32 } type UploadPart struct { @@ -1504,6 +1540,15 @@ type UploadPart struct { Sha256 []byte } +type UploadedMediaReceipt struct { + OwnerUserID int64 + FileID int64 + IntentHash []byte + MediaKind string + MediaID int64 + CreatedAt pgtype.Timestamptz +} + type User struct { ID int64 AccessHash int64 @@ -1637,6 +1682,13 @@ type UserUpdateEvent struct { EventPhone string } +type UserUpdateRetention struct { + UserID int64 + RetainedThroughPts int32 + RetainedThroughDate int32 + UpdatedAt pgtype.Timestamptz +} + type UserUpdateWatermark struct { UserID int64 ContiguousPts int32 diff --git a/internal/store/postgres/sqlcgen/user_update_event.sql.go b/internal/store/postgres/sqlcgen/user_update_event.sql.go index 250e63ce..c864c11c 100644 --- a/internal/store/postgres/sqlcgen/user_update_event.sql.go +++ b/internal/store/postgres/sqlcgen/user_update_event.sql.go @@ -527,29 +527,29 @@ func (q *Queries) BatchListDispatchEvents(ctx context.Context, arg BatchListDisp } const claimDispatchOutbox = `-- name: ClaimDispatchOutbox :many -WITH picked AS ( - SELECT d.target_user_id, d.id - FROM dispatch_outbox d +WITH picked_heads AS ( + SELECT h.target_user_id, h.head_id + FROM dispatch_outbox_user_heads h WHERE ( - d.status = 'pending' - AND d.next_attempt_at <= now() - ) - OR ( - d.status = 'dispatching' - AND d.updated_at < now() - make_interval(secs => $1::int) - ) - ORDER BY d.next_attempt_at ASC, d.target_user_id ASC, d.pts ASC, d.id ASC + h.status = 'pending' + AND h.next_attempt_at <= now() + ) + OR ( + h.status = 'dispatching' + AND h.updated_at < now() - make_interval(secs => $1::int) + ) + ORDER BY h.next_attempt_at ASC, h.target_user_id ASC, h.head_pts ASC, h.head_id ASC LIMIT $2 - FOR UPDATE SKIP LOCKED + FOR UPDATE OF h SKIP LOCKED ) UPDATE dispatch_outbox d SET status = 'dispatching', attempts = d.attempts + 1, updated_at = now() -FROM picked p +FROM picked_heads p WHERE d.target_user_id = p.target_user_id - AND d.id = p.id + AND d.id = p.head_id RETURNING d.id, d.target_user_id, @@ -575,6 +575,8 @@ type ClaimDispatchOutboxRow struct { Attempts int32 } +// durable head 表只保留每用户一行，并同步 head 的 readiness。claim 先锁 +// lane head 再更新对应 outbox 行，既不会扫描 backlog，也不会并发领取同一用户。 func (q *Queries) ClaimDispatchOutbox(ctx context.Context, arg ClaimDispatchOutboxParams) ([]ClaimDispatchOutboxRow, error) { rows, err := q.db.Query(ctx, claimDispatchOutbox, arg.LeaseSeconds, arg.LimitCount) if err != nil { @@ -603,14 +605,100 @@ func (q *Queries) ClaimDispatchOutbox(ctx context.Context, arg ClaimDispatchOutb return items, nil } +const claimDispatchOutboxShards = `-- name: ClaimDispatchOutboxShards :many +WITH picked_heads AS ( + SELECT h.target_user_id, h.head_id + FROM dispatch_outbox_user_heads h + -- 256 与 store.DispatchOutboxLogicalShards、0069 generated column 是同一 + -- schema 常量；不得随 worker 数变化。 + WHERE h.logical_shard = ANY($1::smallint[]) + AND ( + ( + h.status = 'pending' + AND h.next_attempt_at <= now() + ) + OR ( + h.status = 'dispatching' + AND h.updated_at < now() - make_interval(secs => $2::int) + ) + ) + ORDER BY h.next_attempt_at ASC, h.target_user_id ASC, h.head_pts ASC, h.head_id ASC + LIMIT $3 + FOR UPDATE OF h SKIP LOCKED +) +UPDATE dispatch_outbox d +SET + status = 'dispatching', + attempts = d.attempts + 1, + updated_at = now() +FROM picked_heads p +WHERE d.target_user_id = p.target_user_id + AND d.id = p.head_id +RETURNING + d.id, + d.target_user_id, + d.pts, + d.event_type, + d.exclude_auth_key_id, + d.exclude_session_id, + d.attempts +` + +type ClaimDispatchOutboxShardsParams struct { + ShardIds []int16 + LeaseSeconds int32 + LimitCount int32 +} + +type ClaimDispatchOutboxShardsRow struct { + ID int64 + TargetUserID int64 + Pts int32 + EventType string + ExcludeAuthKeyID int64 + ExcludeSessionID int64 + Attempts int32 +} + +// 固定 logical shard 由 target_user_id 决定；运行时 worker 只领取分配给自己的 +// shard 集合，因此同一用户永远只有一条串行 lane，而不同用户可并行。 +func (q *Queries) ClaimDispatchOutboxShards(ctx context.Context, arg ClaimDispatchOutboxShardsParams) ([]ClaimDispatchOutboxShardsRow, error) { + rows, err := q.db.Query(ctx, claimDispatchOutboxShards, arg.ShardIds, arg.LeaseSeconds, arg.LimitCount) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ClaimDispatchOutboxShardsRow + for rows.Next() { + var i ClaimDispatchOutboxShardsRow + if err := rows.Scan( + &i.ID, + &i.TargetUserID, + &i.Pts, + &i.EventType, + &i.ExcludeAuthKeyID, + &i.ExcludeSessionID, + &i.Attempts, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const deleteFailedDispatchOutbox = `-- name: DeleteFailedDispatchOutbox :one -WITH doomed AS ( - SELECT target_user_id, id - FROM dispatch_outbox - WHERE status = 'failed' - AND updated_at < now() - make_interval(secs => $1::int) - ORDER BY updated_at ASC, target_user_id ASC, id ASC +WITH doomed AS MATERIALIZED ( + SELECT h.target_user_id, h.head_id AS id + FROM dispatch_outbox_user_heads h + WHERE h.status = 'failed' + AND h.updated_at < now() - make_interval(secs => $1::int) + ORDER BY h.updated_at ASC, h.target_user_id ASC, h.head_id ASC LIMIT $2 + FOR UPDATE OF h SKIP LOCKED ), deleted AS ( DELETE FROM dispatch_outbox d @@ -628,6 +716,9 @@ type DeleteFailedDispatchOutboxParams struct { LimitCount int32 } +// failed 只能成为 lane head；从 head 表开始并先锁 head，既走 0074 的小索引，也与 +// claim/completion 保持同一 user_heads→outbox 锁序。删除的只是在线任务，durable +// user_update_events 不动，故客户端仍可经 difference 恢复。 func (q *Queries) DeleteFailedDispatchOutbox(ctx context.Context, arg DeleteFailedDispatchOutboxParams) (int32, error) { row := q.db.QueryRow(ctx, deleteFailedDispatchOutbox, arg.OlderThanSeconds, arg.LimitCount) var deleted_count int32 @@ -1087,66 +1178,121 @@ func (q *Queries) ListUserUpdateEventsAfter(ctx context.Context, arg ListUserUpd return items, nil } -const markDispatchDelivered = `-- name: MarkDispatchDelivered :exec -DELETE FROM dispatch_outbox -WHERE target_user_id = $1 - AND id = $2 +const markDispatchDelivered = `-- name: MarkDispatchDelivered :execrows +WITH locked_head AS MATERIALIZED ( + SELECT h.target_user_id + FROM dispatch_outbox_user_heads h + WHERE h.target_user_id = $3::bigint + FOR UPDATE +) +DELETE FROM dispatch_outbox d +USING locked_head h +WHERE d.target_user_id = h.target_user_id + AND d.id = $1::bigint + AND d.status = 'dispatching' + AND d.attempts = $2::int ` type MarkDispatchDeliveredParams struct { - TargetUserID int64 - ID int64 + ID int64 + ExpectedAttempts int32 + TargetUserID int64 } // 方案 A：投递成功即删除。outbox 是任务队列，delivered 行无保留价值 // （消息在 message_boxes、离线补偿在 user_update_events），删除让表维持「未完成任务」小稳态。 -func (q *Queries) MarkDispatchDelivered(ctx context.Context, arg MarkDispatchDeliveredParams) error { - _, err := q.db.Exec(ctx, markDispatchDelivered, arg.TargetUserID, arg.ID) - return err +// claim 的锁序是 user_heads→outbox；completion 必须先显式锁同一 head 再删 outbox， +// 否则租约过期 claim 与完成恰好竞争时会形成 outbox→head / head→outbox 环路。 +func (q *Queries) MarkDispatchDelivered(ctx context.Context, arg MarkDispatchDeliveredParams) (int64, error) { + result, err := q.db.Exec(ctx, markDispatchDelivered, arg.ID, arg.ExpectedAttempts, arg.TargetUserID) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil } -const markDispatchDeliveredBatch = `-- name: MarkDispatchDeliveredBatch :exec +const markDispatchDeliveredBatch = `-- name: MarkDispatchDeliveredBatch :execrows +WITH input AS MATERIALIZED ( + SELECT tu.target_user_id, di.id, ea.attempts + FROM unnest($1::bigint[]) WITH ORDINALITY AS tu(target_user_id, ord) + JOIN unnest($2::bigint[]) WITH ORDINALITY AS di(id, ord) USING (ord) + JOIN unnest($3::int[]) WITH ORDINALITY AS ea(attempts, ord) USING (ord) +), +locked_heads AS MATERIALIZED ( + SELECT h.target_user_id + FROM dispatch_outbox_user_heads h + JOIN (SELECT DISTINCT target_user_id FROM input) i USING (target_user_id) + -- Match ClaimDispatchOutbox[Shards] exactly. A stale-lease claim may lock several + -- dispatching heads while this completion batch locks the same set; a different + -- multi-row order would merely move the deadlock one level up. + ORDER BY h.next_attempt_at, h.target_user_id, h.head_pts, h.head_id + FOR UPDATE OF h +) DELETE FROM dispatch_outbox d -USING unnest($1::bigint[]) WITH ORDINALITY AS tu(target_user_id, ord) -JOIN unnest($2::bigint[]) WITH ORDINALITY AS di(id, ord) USING (ord) -WHERE d.target_user_id = tu.target_user_id - AND d.id = di.id +USING input i, locked_heads h +WHERE d.target_user_id = h.target_user_id + AND d.target_user_id = i.target_user_id + AND d.id = i.id + AND d.status = 'dispatching' + AND d.attempts = i.attempts ` type MarkDispatchDeliveredBatchParams struct { - TargetUserIds []int64 - Ids []int64 + TargetUserIds []int64 + Ids []int64 + ExpectedAttempts []int32 } // 批量删除一批已投递的 (target_user_id, id)；target_user_id 入 WHERE 命中唯一索引并避免串删。 -func (q *Queries) MarkDispatchDeliveredBatch(ctx context.Context, arg MarkDispatchDeliveredBatchParams) error { - _, err := q.db.Exec(ctx, markDispatchDeliveredBatch, arg.TargetUserIds, arg.Ids) - return err +func (q *Queries) MarkDispatchDeliveredBatch(ctx context.Context, arg MarkDispatchDeliveredBatchParams) (int64, error) { + result, err := q.db.Exec(ctx, markDispatchDeliveredBatch, arg.TargetUserIds, arg.Ids, arg.ExpectedAttempts) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil } -const markDispatchFailed = `-- name: MarkDispatchFailed :exec -UPDATE dispatch_outbox +const markDispatchFailed = `-- name: MarkDispatchFailed :execrows +WITH locked_head AS MATERIALIZED ( + SELECT h.target_user_id + FROM dispatch_outbox_user_heads h + WHERE h.target_user_id = $4::bigint + FOR UPDATE +) +UPDATE dispatch_outbox d SET - status = CASE WHEN attempts >= 5 THEN 'failed' ELSE 'pending' END, + status = CASE WHEN d.attempts >= 5 THEN 'failed' ELSE 'pending' END, next_attempt_at = CASE - WHEN attempts >= 5 THEN next_attempt_at - ELSE now() + make_interval(secs => LEAST(60, attempts * attempts)) + WHEN d.attempts >= 5 THEN d.next_attempt_at + ELSE now() + make_interval(secs => LEAST(60, d.attempts * d.attempts)) END, - last_error = $3, + last_error = $1::text, updated_at = now() -WHERE target_user_id = $1 - AND id = $2 +FROM locked_head h +WHERE d.target_user_id = h.target_user_id + AND d.id = $2::bigint + AND d.status = 'dispatching' + AND d.attempts = $3::int ` type MarkDispatchFailedParams struct { - TargetUserID int64 - ID int64 - LastError string + LastError string + ID int64 + ExpectedAttempts int32 + TargetUserID int64 } -func (q *Queries) MarkDispatchFailed(ctx context.Context, arg MarkDispatchFailedParams) error { - _, err := q.db.Exec(ctx, markDispatchFailed, arg.TargetUserID, arg.ID, arg.LastError) - return err +func (q *Queries) MarkDispatchFailed(ctx context.Context, arg MarkDispatchFailedParams) (int64, error) { + result, err := q.db.Exec(ctx, markDispatchFailed, + arg.LastError, + arg.ID, + arg.ExpectedAttempts, + arg.TargetUserID, + ) + if err != nil { + return 0, err + } + return result.RowsAffected(), nil } const maxUserPts = `-- name: MaxUserPts :one diff --git a/internal/store/postgres/update_event_retention.go b/internal/store/postgres/update_event_retention.go new file mode 100644 index 00000000..b214a141 --- /dev/null +++ b/internal/store/postgres/update_event_retention.go @@ -0,0 +1,300 @@ +package postgres + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/jackc/pgx/v5" +) + +const userUpdateRetentionTransactionBatch = 256 + +// DeleteConfirmedPrefix 删除账号 durable update 的共同确认安全前缀。 +// +// 安全边界：只考虑当前 authorizations；任一授权缺 update_states 时其 observed 水位按 0， +// 因而不会删除它可能仍需的事件。AuthorizationStore.Bind 会为新授权以账号当前水位 +// 初始化 delivered state、以已回收 floor 初始化 observed baseline：新设备无需恢复其授权 +// 创建前已删除的事件，但在主动报告后续 pts 前仍会阻塞新的前缀回收。 +func (s *UpdateEventStore) DeleteConfirmedPrefix(ctx context.Context, olderThan time.Duration, limit int) (int, error) { + if s == nil || s.db == nil { + return 0, nil + } + if olderThan <= 0 { + olderThan = 7 * 24 * time.Hour + } + if limit <= 0 { + limit = 10000 + } + if limit > 100000 { + limit = 100000 + } + cutoff := int32(time.Now().Add(-olderThan).Unix()) + deletedTotal := 0 + excluded := make([]int64, 0) + for deletedTotal < limit { + userID, err := s.oldestConfirmedRetentionCandidate(ctx, cutoff, excluded) + if err != nil { + if errors.Is(err, pgx.ErrNoRows) { + break + } + return deletedTotal, err + } + deleted := 0 + err = withTx(ctx, s.db, "delete confirmed user update prefix", func(tx pgx.Tx) error { + var pruneErr error + chunkLimit := limit - deletedTotal + if chunkLimit > userUpdateRetentionTransactionBatch { + chunkLimit = userUpdateRetentionTransactionBatch + } + deleted, pruneErr = pruneConfirmedUserPrefixTx(ctx, tx, userID, cutoff, chunkLimit) + return pruneErr + }) + if err != nil { + return deletedTotal, err + } + if deleted == 0 { + // candidate 在选取与加锁之间可能被其它 worker 处理，或遇到既有空洞； + // 本轮排除后继续找其它用户，避免一个竞态账号饿死全局回收。 + // candidate SQL 已只允许 floor 后的 immediate head，不再让“新 head+旧 tail” + // 或缺口账号占用任意 256-pass 配额；因此这里也不再设人为 256 截断。 + excluded = append(excluded, userID) + continue + } + deletedTotal += deleted + } + return deletedTotal, nil +} + +func (s *UpdateEventStore) oldestConfirmedRetentionCandidate(ctx context.Context, cutoff int32, excluded []int64) (int64, error) { + var userID int64 + err := s.db.QueryRow(ctx, ` +SELECT e.user_id +FROM user_update_events e +LEFT JOIN user_update_retention r ON r.user_id = e.user_id +WHERE e.date < $1 + AND e.pts > COALESCE(r.retained_through_pts, 0) + AND e.pts_count > 0 + -- Only the first complete event immediately after the retained floor may make + -- a user a candidate. A later old-dated tail behind a recent head must not + -- repeatedly win the global date seek and then produce a zero-row prune. + AND e.pts = COALESCE(r.retained_through_pts, 0) + e.pts_count + AND NOT EXISTS ( + SELECT 1 + FROM user_update_events earlier + WHERE earlier.user_id = e.user_id + AND earlier.pts > COALESCE(r.retained_through_pts, 0) + AND earlier.pts < e.pts + ) + AND NOT (e.user_id = ANY($2::bigint[])) + AND EXISTS (SELECT 1 FROM authorizations a WHERE a.user_id = e.user_id) + AND e.pts <= COALESCE(( + SELECT MIN(COALESCE(s.observed_pts, 0)) + FROM authorizations a + LEFT JOIN update_states s + ON s.auth_key_id = a.auth_key_id + AND s.user_id = a.user_id + WHERE a.user_id = e.user_id + ), 0) +ORDER BY e.date ASC, e.user_id ASC, e.pts ASC +LIMIT 1`, cutoff, excluded).Scan(&userID) + if err != nil { + return 0, err + } + return userID, nil +} + +type retainedUserEventRow struct { + pts int + ptsCount int + date int +} + +func pruneConfirmedUserPrefixTx(ctx context.Context, tx pgx.Tx, userID int64, cutoff int32, limit int) (int, error) { + if userID == 0 || limit <= 0 { + return 0, nil + } + // 与所有 pts 分配共享 watermark 行锁：新业务事件不能在 floor 计算与删除之间穿插。 + var currentPts int + if err := tx.QueryRow(ctx, ` +SELECT contiguous_pts +FROM user_update_watermarks +WHERE user_id = $1 +FOR UPDATE`, userID).Scan(¤tPts); err != nil { + if errors.Is(err, pgx.ErrNoRows) { + return 0, nil + } + return 0, fmt.Errorf("lock user update watermark: %w", err) + } + if _, err := tx.Exec(ctx, ` +INSERT INTO user_update_retention (user_id) +VALUES ($1) +ON CONFLICT (user_id) DO NOTHING`, userID); err != nil { + return 0, fmt.Errorf("ensure user update retention: %w", err) + } + var floor int + if err := tx.QueryRow(ctx, ` +SELECT retained_through_pts +FROM user_update_retention +WHERE user_id = $1 +FOR UPDATE`, userID).Scan(&floor); err != nil { + return 0, fmt.Errorf("lock user update retention: %w", err) + } + var authCount, safePts int + if err := tx.QueryRow(ctx, ` +SELECT COUNT(*)::int, COALESCE(MIN(COALESCE(s.observed_pts, 0)), 0)::int +FROM authorizations a +LEFT JOIN update_states s + ON s.auth_key_id = a.auth_key_id + AND s.user_id = a.user_id +WHERE a.user_id = $1`, userID).Scan(&authCount, &safePts); err != nil { + return 0, fmt.Errorf("load confirmed user update floor: %w", err) + } + if authCount == 0 || safePts <= floor { + return 0, nil + } + if safePts > currentPts { + return 0, fmt.Errorf("confirmed user update pts %d exceeds current %d for user %d", safePts, currentPts, userID) + } + rows, err := tx.Query(ctx, ` +SELECT pts, pts_count, date +FROM user_update_events +WHERE user_id = $1 + AND pts > $2 + AND pts <= $3 + AND date < $4 +ORDER BY pts ASC +LIMIT $5`, userID, floor, safePts, cutoff, limit) + if err != nil { + return 0, fmt.Errorf("list confirmed user update prefix: %w", err) + } + defer rows.Close() + events := make([]retainedUserEventRow, 0, limit) + expected := floor + for rows.Next() { + var event retainedUserEventRow + if err := rows.Scan(&event.pts, &event.ptsCount, &event.date); err != nil { + return 0, fmt.Errorf("scan confirmed user update prefix: %w", err) + } + if event.ptsCount <= 0 { + return 0, fmt.Errorf("invalid pts_count %d at user %d pts %d", event.ptsCount, userID, event.pts) + } + expected += event.ptsCount + if event.pts != expected { + // 绝不跨越既有空洞推进 retained floor。 + break + } + events = append(events, event) + } + if err := rows.Err(); err != nil { + rows.Close() + return 0, fmt.Errorf("iterate confirmed user update prefix: %w", err) + } + // A gap/date boundary may stop iteration before pgx consumed the result set. Close explicitly + // before issuing DELETE on the same transaction connection; otherwise pgx reports conn busy. + rows.Close() + if len(events) == 0 { + return 0, nil + } + pts := make([]int32, len(events)) + for i, event := range events { + pts[i] = int32(event.pts) + } + // Every outbox mutation follows user_heads→outbox. Retention may remove a pending or leased + // task after the client has already confirmed its durable event; lock the lane head first so + // it cannot deadlock a lease-expiry claim/completion. No head means these events have no online + // task and the durable prefix can still be pruned safely. + var lockedDispatchUserID int64 + err = tx.QueryRow(ctx, ` +SELECT target_user_id +FROM dispatch_outbox_user_heads +WHERE target_user_id = $1 +FOR UPDATE`, userID).Scan(&lockedDispatchUserID) + if err != nil && !errors.Is(err, pgx.ErrNoRows) { + return 0, fmt.Errorf("lock retained user update dispatch head: %w", err) + } + // A retained durable event can still have a pending/dispatching outbox row (for example a + // client confirmed the pts through difference while an online push lease was in flight). + // Remove those leases first, in this same transaction. The outbox head trigger promotes the + // next user lane row; any worker holding an old attempts token is fenced by MarkDelivered/ + // MarkFailed returning ErrDispatchLeaseLost after this commit. + if _, err := tx.Exec(ctx, ` +DELETE FROM dispatch_outbox +WHERE target_user_id = $1 + AND pts = ANY($2::int[])`, userID, pts); err != nil { + return 0, fmt.Errorf("delete retained user update dispatch outbox: %w", err) + } + tag, err := tx.Exec(ctx, ` +DELETE FROM user_update_events +WHERE user_id = $1 + AND pts = ANY($2::int[])`, userID, pts) + if err != nil { + return 0, fmt.Errorf("delete confirmed user update prefix: %w", err) + } + if tag.RowsAffected() != int64(len(events)) { + return 0, fmt.Errorf("delete confirmed user update prefix affected %d rows, want %d", tag.RowsAffected(), len(events)) + } + last := events[len(events)-1] + if _, err := tx.Exec(ctx, ` +UPDATE user_update_retention +SET retained_through_pts = $2, + retained_through_date = $3, + updated_at = now() +WHERE user_id = $1`, userID, last.pts, last.date); err != nil { + return 0, fmt.Errorf("advance user update retention: %w", err) + } + return len(events), nil +} + +// UserUpdateRetentionCheckpoint 返回当前 auth key 已明确确认、可通过普通 +// differenceSlice 跳过的安全前缀。 +// +// 一旦 retained floor > 0，仍存在的 authorization 必须同时有 observed_pts >= floor； +// AuthorizationStore.Bind 在同一事务建立这个 baseline。若这里看到 authorization 存在但 +// state 缺失/倒退，说明生命周期不变量已经破坏。此时必须 fail-fast，不能返回 ok=false 后让 +// GetDifference 从已删除前缀继续读并伪装成空差分。 +func (s *UpdateEventStore) UserUpdateRetentionCheckpoint(ctx context.Context, authKeyID [8]byte, userID int64) (pts, date int, ok bool, err error) { + if s == nil || s.db == nil || userID == 0 || authKeyID == ([8]byte{}) { + return 0, 0, false, nil + } + var ( + authorized bool + observed int + ) + err = s.db.QueryRow(ctx, ` +SELECT + r.retained_through_pts, + r.retained_through_date, + EXISTS ( + SELECT 1 + FROM authorizations a + WHERE a.auth_key_id = $1 + AND a.user_id = r.user_id + ) AS authorized, + COALESCE(( + SELECT s.observed_pts + FROM update_states s + WHERE s.auth_key_id = $1 + AND s.user_id = r.user_id + ), -1)::int AS observed_pts +FROM user_update_retention r +WHERE r.user_id = $2 + AND r.retained_through_pts > 0`, authKeyIDToInt64(authKeyID), userID).Scan(&pts, &date, &authorized, &observed) + if errors.Is(err, pgx.ErrNoRows) { + return 0, 0, false, nil + } + if err != nil { + return 0, 0, false, fmt.Errorf("get user update retention checkpoint: %w", err) + } + if !authorized { + return 0, 0, false, nil + } + if observed < pts { + return 0, 0, false, fmt.Errorf( + "get user update retention checkpoint: invariant violation: auth key %x user %d observed pts %d below retained floor %d", + authKeyID, userID, observed, pts, + ) + } + return pts, date, true, nil +} diff --git a/internal/store/postgres/update_event_retention_integration_test.go b/internal/store/postgres/update_event_retention_integration_test.go new file mode 100644 index 00000000..67d79828 --- /dev/null +++ b/internal/store/postgres/update_event_retention_integration_test.go @@ -0,0 +1,643 @@ +package postgres + +import ( + "context" + "crypto/rand" + "errors" + "fmt" + "testing" + "time" + + appupdates "telesrv/internal/app/updates" + "telesrv/internal/domain" + "telesrv/internal/store" +) + +func TestUserUpdateRetentionUsesClientObservedCommonPrefixPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + userID := createRevokeTestUser(t, ctx, pool, "update-retention") + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + states := NewUpdateStateStore(pool) + events := NewUpdateEventStore(pool) + + newKey := func() [8]byte { + var id [8]byte + if _, err := rand.Read(id[:]); err != nil { + t.Fatalf("random auth key id: %v", err) + } + return id + } + authOne, authTwo := newKey(), newKey() + for _, id := range [][8]byte{authOne, authTwo} { + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key %x: %v", id, err) + } + id := id + t.Cleanup(func() { _ = keys.Delete(ctx, id) }) + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: id, UserID: userID}); err != nil { + t.Fatalf("bind authorization %x: %v", id, err) + } + } + + const oldDate = 1_600_000_000 + for i := 1; i <= 3; i++ { + if _, err := events.AppendAllocated(ctx, userID, domain.UpdateEvent{ + Type: domain.UpdateEventNoop, PtsCount: 1, Date: oldDate + i, + }); err != nil { + t.Fatalf("append event %d: %v", i, err) + } + } + + // Save is the state the server has sent/constructed. Neither device has proved receipt, so it + // must not authorize retention even though both delivered cursors are at pts=3. + for _, id := range [][8]byte{authOne, authTwo} { + if err := states.Save(ctx, id, userID, domain.UpdateState{Pts: 3, Date: oldDate + 3}); err != nil { + t.Fatalf("save delivered state %x: %v", id, err) + } + } + if deleted, err := events.DeleteConfirmedPrefix(ctx, time.Second, 10); err != nil || deleted != 0 { + t.Fatalf("delete with no observed cursor = %d/%v, want 0/nil", deleted, err) + } + + if err := states.ObserveClientState(ctx, authOne, userID, domain.UpdateState{Pts: 3, Date: oldDate + 3}); err != nil { + t.Fatalf("observe first device: %v", err) + } + if deleted, err := events.DeleteConfirmedPrefix(ctx, time.Second, 10); err != nil || deleted != 0 { + t.Fatalf("delete while second device unobserved = %d/%v, want 0/nil", deleted, err) + } + + // Common observed floor=min(3,1)=1, so exactly the first contiguous event is removable. + if err := states.ObserveClientState(ctx, authTwo, userID, domain.UpdateState{Pts: 1, Date: oldDate + 1}); err != nil { + t.Fatalf("observe second device pts=1: %v", err) + } + deleted, err := events.DeleteConfirmedPrefix(ctx, time.Second, 10) + if err != nil || deleted != 1 { + t.Fatalf("delete common prefix = %d/%v, want 1/nil", deleted, err) + } + pts, date, ok, err := events.UserUpdateRetentionCheckpoint(ctx, authTwo, userID) + if err != nil || !ok || pts != 1 || date != oldDate+1 { + t.Fatalf("checkpoint = pts:%d date:%d ok:%v err:%v, want 1/%d/true/nil", pts, date, ok, err, oldDate+1) + } + remaining, err := events.ListAfter(ctx, userID, 0, 10) + if err != nil || len(remaining) != 2 || remaining[0].Pts != 2 || remaining[1].Pts != 3 { + t.Fatalf("remaining events = %+v err=%v, want pts 2,3", remaining, err) + } + + if err := states.ObserveClientState(ctx, authTwo, userID, domain.UpdateState{Pts: 3, Date: oldDate + 3}); err != nil { + t.Fatalf("observe second device pts=3: %v", err) + } + deleted, err = events.DeleteConfirmedPrefix(ctx, time.Second, 10) + if err != nil || deleted != 2 { + t.Fatalf("delete remaining common prefix = %d/%v, want 2/nil", deleted, err) + } + + // A newly created authorization did not exist when the common prefix was confirmed. Seed its + // observed baseline at the retained floor (not at current pts): it can receive an ordinary + // empty differenceSlice checkpoint instead of falling into a silent hole, while still blocking + // any future pruning until it reports subsequent progress itself. + authThree := newKey() + if err := keys.Save(ctx, store.AuthKeyData{ID: authThree}); err != nil { + t.Fatalf("save third auth key: %v", err) + } + t.Cleanup(func() { _ = keys.Delete(ctx, authThree) }) + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: authThree, UserID: userID}); err != nil { + t.Fatalf("bind third authorization: %v", err) + } + pts, date, ok, err = events.UserUpdateRetentionCheckpoint(ctx, authThree, userID) + if err != nil || !ok || pts != 3 || date != oldDate+3 { + t.Fatalf("new authorization checkpoint = pts:%d date:%d ok:%v err:%v, want 3/%d/true/nil", pts, date, ok, err, oldDate+3) + } + var observed int + if err := pool.QueryRow(ctx, ` +SELECT observed_pts FROM update_states WHERE auth_key_id = $1 AND user_id = $2 +`, authKeyIDToInt64(authThree), userID).Scan(&observed); err != nil { + t.Fatalf("load third observed floor: %v", err) + } + if observed != 3 { + t.Fatalf("new authorization observed_pts = %d, want retained floor 3", observed) + } +} + +func TestAuthorizationBindSwitchesAccountAfterRetainedFloorPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + oldUserID := createRevokeTestUser(t, ctx, pool, "retention-switch-old") + newUserID := createRevokeTestUser(t, ctx, pool, "retention-switch-new") + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + states := NewUpdateStateStore(pool) + events := NewUpdateEventStore(pool) + mainKey := randomUpdateRetentionAuthKey(t) + guardKey := randomUpdateRetentionAuthKey(t) + for _, id := range [][8]byte{mainKey, guardKey} { + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key %x: %v", id, err) + } + } + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM update_states WHERE auth_key_id = ANY($1::bigint[])", []int64{ + authKeyIDToInt64(mainKey), authKeyIDToInt64(guardKey), + }) + _ = keys.Delete(ctx, mainKey) + _ = keys.Delete(ctx, guardKey) + }) + + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: mainKey, UserID: oldUserID}); err != nil { + t.Fatalf("bind main key to old account: %v", err) + } + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: guardKey, UserID: newUserID}); err != nil { + t.Fatalf("bind guard key to new account: %v", err) + } + const oldDate = 1_600_100_000 + for i := 1; i <= 3; i++ { + if _, err := events.AppendAllocated(ctx, newUserID, domain.UpdateEvent{ + Type: domain.UpdateEventNoop, PtsCount: 1, Date: oldDate + i, + }); err != nil { + t.Fatalf("append new-account event %d: %v", i, err) + } + } + if err := states.ObserveClientState(ctx, guardKey, newUserID, domain.UpdateState{Pts: 3, Date: oldDate + 3}); err != nil { + t.Fatalf("observe guard through pts 3: %v", err) + } + if deleted, err := events.DeleteConfirmedPrefix(ctx, time.Second, 10); err != nil || deleted != 3 { + t.Fatalf("prune new-account prefix = %d/%v, want 3/nil", deleted, err) + } + + // This is the account-switch boundary that used to be followed by Router.ClearAuthKey, + // deleting the state Bind had just created for newUserID. + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: mainKey, UserID: newUserID}); err != nil { + t.Fatalf("switch main key to new account: %v", err) + } + var oldStates int + if err := pool.QueryRow(ctx, ` +SELECT COUNT(*)::int +FROM update_states +WHERE auth_key_id = $1 AND user_id = $2`, authKeyIDToInt64(mainKey), oldUserID).Scan(&oldStates); err != nil { + t.Fatalf("count old-account states: %v", err) + } + if oldStates != 0 { + t.Fatalf("old-account update states = %d, want 0", oldStates) + } + var delivered, observed int + if err := pool.QueryRow(ctx, ` +SELECT pts, observed_pts +FROM update_states +WHERE auth_key_id = $1 AND user_id = $2`, authKeyIDToInt64(mainKey), newUserID).Scan(&delivered, &observed); err != nil { + t.Fatalf("load switched-account state: %v", err) + } + if delivered != 3 || observed != 3 { + t.Fatalf("switched-account state = delivered:%d observed:%d, want 3/3", delivered, observed) + } + + diff, err := appupdates.NewService(states, events).GetDifference( + ctx, + mainKey, + newUserID, + domain.UpdateState{Pts: 0, Date: oldDate}, + ) + if err != nil { + t.Fatalf("difference after account switch: %v", err) + } + if !diff.Partial || len(diff.Events) != 0 || diff.State.Pts != 3 { + t.Fatalf("switch checkpoint difference = %+v, want empty slice at retained pts 3", diff) + } +} + +func TestAuthorizationBindRejectsFutureSameUserStatePostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + userID := createRevokeTestUser(t, ctx, pool, "retention-stale-rebind") + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + states := NewUpdateStateStore(pool) + events := NewUpdateEventStore(pool) + guardKey := randomUpdateRetentionAuthKey(t) + staleKey := randomUpdateRetentionAuthKey(t) + for _, id := range [][8]byte{guardKey, staleKey} { + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key %x: %v", id, err) + } + } + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM update_states WHERE auth_key_id = ANY($1::bigint[])", []int64{ + authKeyIDToInt64(guardKey), authKeyIDToInt64(staleKey), + }) + _ = keys.Delete(ctx, guardKey) + _ = keys.Delete(ctx, staleKey) + }) + + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: guardKey, UserID: userID}); err != nil { + t.Fatalf("bind guard authorization: %v", err) + } + const oldDate = 1_600_200_000 + for i := 1; i <= 5; i++ { + if _, err := events.AppendAllocated(ctx, userID, domain.UpdateEvent{ + Type: domain.UpdateEventNoop, PtsCount: 1, Date: oldDate + i, + }); err != nil { + t.Fatalf("append event %d: %v", i, err) + } + } + if err := states.ObserveClientState(ctx, guardKey, userID, domain.UpdateState{Pts: 3, Date: oldDate + 3}); err != nil { + t.Fatalf("observe guard pts 3: %v", err) + } + if deleted, err := events.DeleteConfirmedPrefix(ctx, time.Second, 10); err != nil || deleted != 3 { + t.Fatalf("prune stale-rebind prefix = %d/%v, want 3/nil", deleted, err) + } + + // Deliberately inject historical corruption: the authorization is absent while a stale cursor + // claims a future pts beyond the account's contiguous watermark (5). Bind must fail-fast and + // leave the key unauthorized; preserving pts=7 would make future retention/difference lie. + if _, err := pool.Exec(ctx, ` +INSERT INTO update_states (auth_key_id, user_id, pts, qts, date, seq, observed_pts) +VALUES ($1, $2, 7, 4, $3, 2, 1)`, authKeyIDToInt64(staleKey), userID, oldDate+1); err != nil { + t.Fatalf("insert stale update state: %v", err) + } + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: staleKey, UserID: userID}); err == nil { + t.Fatal("Bind accepted future update state, want invariant error") + } + if _, found, err := auths.ByAuthKey(ctx, staleKey); err != nil || found { + t.Fatalf("authorization after rejected Bind found=%v err=%v, want false/nil", found, err) + } + var delivered, qts, seq, observed int + if err := pool.QueryRow(ctx, ` +SELECT pts, qts, seq, observed_pts +FROM update_states +WHERE auth_key_id = $1 AND user_id = $2`, authKeyIDToInt64(staleKey), userID).Scan(&delivered, &qts, &seq, &observed); err != nil { + t.Fatalf("load rejected stale state: %v", err) + } + if delivered != 7 || qts != 4 || seq != 2 || observed != 1 { + t.Fatalf("rejected stale state mutated = pts:%d qts:%d seq:%d observed:%d, want 7/4/2/1", delivered, qts, seq, observed) + } + + // Once an explicit repair brings the persisted cursor back inside the current account + // watermark, Bind may establish the retained-floor baseline without moving qts/seq backwards. + if _, err := pool.Exec(ctx, ` +UPDATE update_states +SET pts = 5 +WHERE auth_key_id = $1 AND user_id = $2`, authKeyIDToInt64(staleKey), userID); err != nil { + t.Fatalf("repair future delivered state: %v", err) + } + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: staleKey, UserID: userID}); err != nil { + t.Fatalf("bind explicitly repaired authorization: %v", err) + } + if err := pool.QueryRow(ctx, ` +SELECT pts, qts, seq, observed_pts +FROM update_states +WHERE auth_key_id = $1 AND user_id = $2`, authKeyIDToInt64(staleKey), userID).Scan(&delivered, &qts, &seq, &observed); err != nil { + t.Fatalf("load bound repaired state: %v", err) + } + if delivered != 5 || qts != 4 || seq != 2 || observed != 3 { + t.Fatalf("bound repaired state = pts:%d qts:%d seq:%d observed:%d, want 5/4/2/3", delivered, qts, seq, observed) + } + + // Protection: if the lifecycle invariant is corrupted again, checkpoint lookup must fail-fast + // instead of letting GetDifference fall through to an empty read below deleted history. + if _, err := pool.Exec(ctx, ` +UPDATE update_states +SET observed_pts = 1 +WHERE auth_key_id = $1 AND user_id = $2`, authKeyIDToInt64(staleKey), userID); err != nil { + t.Fatalf("corrupt observed state for guard test: %v", err) + } + if _, _, _, err := events.UserUpdateRetentionCheckpoint(ctx, staleKey, userID); err == nil { + t.Fatal("checkpoint with observed below retained floor succeeded, want invariant error") + } + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: staleKey, UserID: userID}); err != nil { + t.Fatalf("same-user Bind did not repair observed floor: %v", err) + } + if pts, _, ok, err := events.UserUpdateRetentionCheckpoint(ctx, staleKey, userID); err != nil || !ok || pts != 3 { + t.Fatalf("checkpoint after same-user repair = pts:%d ok:%v err:%v", pts, ok, err) + } +} + +func TestAuthorizationBindSerializesWithRetentionTwoConnectionsPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + userID := createRevokeTestUser(t, ctx, pool, "retention-bind-race") + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + states := NewUpdateStateStore(pool) + events := NewUpdateEventStore(pool) + guardKey := randomUpdateRetentionAuthKey(t) + newKey := randomUpdateRetentionAuthKey(t) + for _, id := range [][8]byte{guardKey, newKey} { + if err := keys.Save(ctx, store.AuthKeyData{ID: id}); err != nil { + t.Fatalf("save auth key %x: %v", id, err) + } + id := id + t.Cleanup(func() { _ = keys.Delete(ctx, id) }) + } + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: guardKey, UserID: userID}); err != nil { + t.Fatalf("bind guard authorization: %v", err) + } + const eventDate = 1_600_300_001 + if _, err := events.AppendAllocated(ctx, userID, domain.UpdateEvent{ + Type: domain.UpdateEventNoop, PtsCount: 1, Date: eventDate, + }); err != nil { + t.Fatalf("append guarded event: %v", err) + } + if err := states.ObserveClientState(ctx, guardKey, userID, domain.UpdateState{Pts: 1, Date: eventDate}); err != nil { + t.Fatalf("observe guard event: %v", err) + } + + retentionConn, err := pool.Acquire(ctx) + if err != nil { + t.Fatalf("acquire retention connection: %v", err) + } + defer retentionConn.Release() + bindConn, err := pool.Acquire(ctx) + if err != nil { + t.Fatalf("acquire bind connection: %v", err) + } + defer bindConn.Release() + + tx, err := retentionConn.Begin(ctx) + if err != nil { + t.Fatalf("begin retention transaction: %v", err) + } + committed := false + defer func() { + if !committed { + _ = tx.Rollback(context.Background()) + } + }() + var currentPts, floor int + if err := tx.QueryRow(ctx, ` +SELECT contiguous_pts +FROM user_update_watermarks +WHERE user_id = $1 +FOR UPDATE`, userID).Scan(¤tPts); err != nil { + t.Fatalf("lock retention watermark: %v", err) + } + if err := tx.QueryRow(ctx, ` +SELECT retained_through_pts +FROM user_update_retention +WHERE user_id = $1 +FOR UPDATE`, userID).Scan(&floor); err != nil { + t.Fatalf("lock retention floor: %v", err) + } + if currentPts != 1 || floor != 0 { + t.Fatalf("pre-race watermark/floor = %d/%d, want 1/0", currentPts, floor) + } + + bindCtx, cancelBind := context.WithTimeout(ctx, 5*time.Second) + defer cancelBind() + bindDone := make(chan error, 1) + go func() { + bindDone <- NewAuthorizationStore(bindConn).Bind(bindCtx, domain.Authorization{ + AuthKeyID: newKey, + UserID: userID, + }) + }() + + // Observe the second physical connection waiting on the watermark row. This proves the + // synchronization is a database lock, rather than relying on scheduler timing in the test. + bindPID := bindConn.Conn().PgConn().PID() + waitDeadline := time.Now().Add(2 * time.Second) + waiting := false + for time.Now().Before(waitDeadline) { + select { + case err := <-bindDone: + t.Fatalf("Bind completed before retained-floor transaction committed: %v", err) + default: + } + if err := tx.QueryRow(ctx, ` +SELECT COALESCE(wait_event_type = 'Lock', false) +FROM pg_stat_activity +WHERE pid = $1`, bindPID).Scan(&waiting); err != nil { + t.Fatalf("inspect bind lock wait: %v", err) + } + if waiting { + break + } + time.Sleep(10 * time.Millisecond) + } + if !waiting { + t.Fatal("Bind connection did not wait on retention watermark lock") + } + + // Complete the valid confirmed-prefix transition while Bind is waiting. After commit Bind + // must read floor=1 and atomically seed observed_pts=1; floor=0 would create a silent hole. + if tag, err := tx.Exec(ctx, ` +DELETE FROM user_update_events +WHERE user_id = $1 AND pts = 1`, userID); err != nil || tag.RowsAffected() != 1 { + t.Fatalf("delete retained event rows=%d err=%v, want 1/nil", tag.RowsAffected(), err) + } + if _, err := tx.Exec(ctx, ` +UPDATE user_update_retention +SET retained_through_pts = 1, + retained_through_date = $2, + updated_at = now() +WHERE user_id = $1`, userID, eventDate); err != nil { + t.Fatalf("advance retained floor: %v", err) + } + if err := tx.Commit(ctx); err != nil { + t.Fatalf("commit retained floor: %v", err) + } + committed = true + select { + case err := <-bindDone: + if err != nil { + t.Fatalf("Bind after retention commit: %v", err) + } + case <-time.After(5 * time.Second): + t.Fatal("Bind remained blocked after retention commit") + } + + var delivered, observed int + if err := pool.QueryRow(ctx, ` +SELECT pts, observed_pts +FROM update_states +WHERE auth_key_id = $1 AND user_id = $2`, authKeyIDToInt64(newKey), userID).Scan(&delivered, &observed); err != nil { + t.Fatalf("load raced bind baseline: %v", err) + } + if delivered != 1 || observed != 1 { + t.Fatalf("raced bind baseline = delivered:%d observed:%d, want 1/1", delivered, observed) + } +} + +func TestUserUpdateRetentionOldTailsDoNotConsumeCandidatePassPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + const tailUsers = 256 + const totalUsers = tailUsers + 1 + prefix := fmt.Sprintf("+188%010d", time.Now().UnixNano()%10_000_000_000) + rows, err := pool.Query(ctx, ` +INSERT INTO users (access_hash, phone, first_name) +SELECT $1::bigint + n, $2 || lpad(n::text, 3, '0'), 'retention-old-tail' +FROM generate_series(1, $3::int) AS n +RETURNING id +`, time.Now().UnixNano(), prefix, totalUsers) + if err != nil { + t.Fatalf("bulk insert old-tail users: %v", err) + } + userIDs := make([]int64, 0, totalUsers) + for rows.Next() { + var userID int64 + if err := rows.Scan(&userID); err != nil { + rows.Close() + t.Fatalf("scan old-tail user: %v", err) + } + userIDs = append(userIDs, userID) + } + if err := rows.Err(); err != nil { + rows.Close() + t.Fatalf("iterate old-tail users: %v", err) + } + rows.Close() + if len(userIDs) != totalUsers { + t.Fatalf("inserted users = %d, want %d", len(userIDs), totalUsers) + } + authKeyIDs := make([]int64, len(userIDs)) + watermarks := make([]int32, len(userIDs)) + for i, userID := range userIDs { + authKeyIDs[i] = -userID + if i < tailUsers { + watermarks[i] = 2 + } else { + watermarks[i] = 1 + } + } + t.Cleanup(func() { + _, _ = pool.Exec(ctx, "DELETE FROM update_states WHERE auth_key_id = ANY($1::bigint[])", authKeyIDs) + _, _ = pool.Exec(ctx, "DELETE FROM auth_keys WHERE auth_key_id = ANY($1::bigint[])", authKeyIDs) + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", userIDs) + }) + if _, err := pool.Exec(ctx, ` +INSERT INTO auth_keys (auth_key_id, body, server_salt) +SELECT id, decode(repeat('00', 256), 'hex'), 0 +FROM unnest($1::bigint[]) AS id`, authKeyIDs); err != nil { + t.Fatalf("bulk insert old-tail auth keys: %v", err) + } + if _, err := pool.Exec(ctx, ` +INSERT INTO authorizations (auth_key_id, user_id) +SELECT * FROM unnest($1::bigint[], $2::bigint[])`, authKeyIDs, userIDs); err != nil { + t.Fatalf("bulk insert old-tail authorizations: %v", err) + } + if _, err := pool.Exec(ctx, ` +INSERT INTO user_update_watermarks (user_id, contiguous_pts) +SELECT * FROM unnest($1::bigint[], $2::integer[])`, userIDs, watermarks); err != nil { + t.Fatalf("bulk insert old-tail watermarks: %v", err) + } + if _, err := pool.Exec(ctx, ` +INSERT INTO update_states (auth_key_id, user_id, pts, observed_pts) +SELECT auth_key_id, user_id, pts, pts +FROM unnest($1::bigint[], $2::bigint[], $3::integer[]) AS input(auth_key_id, user_id, pts)`, authKeyIDs, userIDs, watermarks); err != nil { + t.Fatalf("bulk insert old-tail states: %v", err) + } + recentHeadDate := int32(time.Now().Add(time.Hour).Unix()) + if _, err := pool.Exec(ctx, ` +INSERT INTO user_update_events (user_id, pts, pts_count, date, event_type) +SELECT user_id, 1, 1, $2, 'noop' +FROM unnest($1::bigint[]) AS user_id`, userIDs[:tailUsers], recentHeadDate); err != nil { + t.Fatalf("insert recent old-tail heads: %v", err) + } + if _, err := pool.Exec(ctx, ` +INSERT INTO user_update_events (user_id, pts, pts_count, date, event_type) +SELECT user_id, 2, 1, 1, 'noop' +FROM unnest($1::bigint[]) AS user_id`, userIDs[:tailUsers]); err != nil { + t.Fatalf("insert old tails: %v", err) + } + healthyUserID := userIDs[len(userIDs)-1] + if _, err := pool.Exec(ctx, ` +INSERT INTO user_update_events (user_id, pts, pts_count, date, event_type) +VALUES ($1, 1, 1, 2, 'noop')`, healthyUserID); err != nil { + t.Fatalf("insert healthy retention head: %v", err) + } + + deleted, err := NewUpdateEventStore(pool).DeleteConfirmedPrefix(ctx, time.Second, 1) + if err != nil || deleted != 1 { + t.Fatalf("delete after 256 old tails = %d/%v, want healthy 1/nil", deleted, err) + } + var healthyRows, tailRows int + if err := pool.QueryRow(ctx, ` +SELECT + (SELECT count(*) FROM user_update_events WHERE user_id = $1)::int, + (SELECT count(*) FROM user_update_events WHERE user_id = ANY($2::bigint[]))::int`, healthyUserID, userIDs[:tailUsers]).Scan(&healthyRows, &tailRows); err != nil { + t.Fatalf("count old-tail retention rows: %v", err) + } + if healthyRows != 0 || tailRows != tailUsers*2 { + t.Fatalf("remaining healthy/tail rows = %d/%d, want 0/%d", healthyRows, tailRows, tailUsers*2) + } +} + +func TestUserUpdateRetentionDeletesDispatchLeaseAndPromotesHeadPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + userID := createRevokeTestUser(t, ctx, pool, "retention-dispatch-lease") + keys := NewAuthKeyStore(pool) + auths := NewAuthorizationStore(pool) + states := NewUpdateStateStore(pool) + events := NewUpdateEventStore(pool) + outbox := NewDispatchOutboxStore(pool, WithLeaseTimeout(time.Hour)) + authKeyID := randomUpdateRetentionAuthKey(t) + if err := keys.Save(ctx, store.AuthKeyData{ID: authKeyID}); err != nil { + t.Fatalf("save retention dispatch auth key: %v", err) + } + t.Cleanup(func() { _ = keys.Delete(ctx, authKeyID) }) + if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: authKeyID, UserID: userID}); err != nil { + t.Fatalf("bind retention dispatch authorization: %v", err) + } + appendDispatch := func(date int) domain.UpdateEvent { + t.Helper() + event, err := events.AppendAllocatedWithDispatch(ctx, userID, domain.UpdateEvent{ + Type: domain.UpdateEventDialogPinned, + PtsCount: 1, + Date: date, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: userID}, + Bool: true, + }, [8]byte{}, 0) + if err != nil { + t.Fatalf("append retention dispatch event: %v", err) + } + return event + } + first := appendDispatch(1) + second := appendDispatch(2) + claimed := store.DispatchOutboxItem{TargetUserID: userID, Pts: first.Pts} + if err := pool.QueryRow(ctx, ` +UPDATE dispatch_outbox +SET status = 'dispatching', + attempts = attempts + 1, + updated_at = now() +WHERE target_user_id = $1 AND pts = $2 +RETURNING id, attempts`, userID, first.Pts).Scan(&claimed.ID, &claimed.Attempts); err != nil { + t.Fatalf("acquire exact retention dispatch lease: %v", err) + } + if err := states.ObserveClientState(ctx, authKeyID, userID, domain.UpdateState{Pts: first.Pts, Date: first.Date}); err != nil { + t.Fatalf("observe retained dispatch pts: %v", err) + } + deleted, err := events.DeleteConfirmedPrefix(ctx, time.Second, 1) + if err != nil || deleted != 1 { + t.Fatalf("delete retained dispatch prefix = %d/%v, want 1/nil", deleted, err) + } + + // The in-flight worker owns an attempts token for a row retention just removed. It must be + // fenced instead of recreating/marking the deleted head, while the next pts becomes claimable. + if err := outbox.MarkDelivered(ctx, claimed); !errors.Is(err, store.ErrDispatchLeaseLost) { + t.Fatalf("deliver retained dispatch lease err = %v, want ErrDispatchLeaseLost", err) + } + var eventRows, outboxRows, headPts int + var headStatus string + if err := pool.QueryRow(ctx, ` +SELECT + (SELECT count(*) FROM user_update_events WHERE user_id = $1 AND pts = $2)::int, + (SELECT count(*) FROM dispatch_outbox WHERE target_user_id = $1 AND pts = $2)::int, + (SELECT head_pts FROM dispatch_outbox_user_heads WHERE target_user_id = $1), + (SELECT status FROM dispatch_outbox_user_heads WHERE target_user_id = $1)`, userID, first.Pts).Scan(&eventRows, &outboxRows, &headPts, &headStatus); err != nil { + t.Fatalf("load retained dispatch/head state: %v", err) + } + if eventRows != 0 || outboxRows != 0 || headPts != second.Pts || headStatus != "pending" { + t.Fatalf("retained event/outbox/head = %d/%d/%d/%s, want 0/0/%d/pending", eventRows, outboxRows, headPts, headStatus, second.Pts) + } +} + +func randomUpdateRetentionAuthKey(t *testing.T) [8]byte { + t.Helper() + var id [8]byte + if _, err := rand.Read(id[:]); err != nil { + t.Fatalf("random auth key id: %v", err) + } + return id +} diff --git a/internal/store/postgres/updatestate.go b/internal/store/postgres/updatestate.go index e8213fb3..2d486fd1 100644 --- a/internal/store/postgres/updatestate.go +++ b/internal/store/postgres/updatestate.go @@ -13,12 +13,13 @@ import ( // UpdateStateStore 用 PostgreSQL 实现 store.UpdateStateStore。 type UpdateStateStore struct { - q *sqlcgen.Queries + q *sqlcgen.Queries + db sqlcgen.DBTX } // NewUpdateStateStore 基于 pgx 连接池（或事务）创建 UpdateStateStore。 func NewUpdateStateStore(db sqlcgen.DBTX) *UpdateStateStore { - return &UpdateStateStore{q: sqlcgen.New(db)} + return &UpdateStateStore{q: sqlcgen.New(db), db: db} } func (s *UpdateStateStore) Get(ctx context.Context, id [8]byte, userID int64) (domain.UpdateState, bool, error) { @@ -54,6 +55,21 @@ func (s *UpdateStateStore) Save(ctx context.Context, id [8]byte, userID int64, s return nil } +func (s *UpdateStateStore) ObserveClientState(ctx context.Context, id [8]byte, userID int64, st domain.UpdateState) error { + if st.Pts < 0 { + st.Pts = 0 + } + if _, err := s.db.Exec(ctx, ` +INSERT INTO update_states (auth_key_id, user_id, pts, qts, date, seq, observed_pts) +VALUES ($1, $2, $3, $4, $5, $6, $3) +ON CONFLICT (auth_key_id, user_id) DO UPDATE SET + observed_pts = GREATEST(update_states.observed_pts, EXCLUDED.observed_pts), + updated_at = now()`, authKeyIDToInt64(id), userID, st.Pts, st.Qts, st.Date, st.Seq); err != nil { + return fmt.Errorf("observe client update state: %w", err) + } + return nil +} + func (s *UpdateStateStore) Delete(ctx context.Context, id [8]byte, userID int64) error { if err := s.q.DeleteUpdateState(ctx, sqlcgen.DeleteUpdateStateParams{ AuthKeyID: authKeyIDToInt64(id), diff --git a/internal/store/private_send_idempotency.go b/internal/store/private_send_idempotency.go new file mode 100644 index 00000000..f737cc13 --- /dev/null +++ b/internal/store/private_send_idempotency.go @@ -0,0 +1,243 @@ +package store + +import ( + "bytes" + "crypto/sha256" + "encoding/json" + "fmt" + + "telesrv/internal/domain" +) + +const privateSendFingerprintVersion = 1 + +const channelSendFingerprintVersion = 1 + +const sendSnapshotVersion = 1 + +type privateSendSnapshotEnvelope struct { + Version int `json:"version"` + Message domain.Message `json:"message"` +} + +type channelSendSnapshotEnvelope struct { + Version int `json:"version"` + Message domain.ChannelMessage `json:"message"` +} + +// privateSendFingerprintPayload 只包含一次发送的客户端不可变意图。Date、origin +// auth/session、当前 block 状态与 automation 元数据均为执行环境，不得让同一请求在 +// 重连或状态变化后变成另一条逻辑消息。sender/random_id 由幂等索引键单独约束。 +type privateSendFingerprintPayload struct { + Version int `json:"version"` + RecipientUserID int64 `json:"recipient_user_id"` + Message string `json:"message"` + Entities []domain.MessageEntity `json:"entities"` + Media *domain.MessageMedia `json:"media"` + Silent bool `json:"silent"` + NoForwards bool `json:"noforwards"` + ReplyTo *domain.MessageReply `json:"reply_to,omitempty"` + Forward *domain.MessageForward `json:"forward,omitempty"` + TTLPeriod int `json:"ttl_period"` + ViaBotID int64 `json:"via_bot_id"` + GroupedID int64 `json:"grouped_id"` + Effect int64 `json:"effect"` + ReplyMarkup *domain.MessageReplyMarkup `json:"reply_markup"` + RichMessage *domain.MessageRichMessage `json:"rich_message"` +} + +// channelSendFingerprintPayload contains only the durable intent of an internal channel send. +// Operational projection fields (Date, recipient/mention lists, PostAuthor and lookup hints) are +// intentionally absent: they may change after the first commit and cannot redefine a replay. +type channelSendFingerprintPayload struct { + Version int `json:"version"` + ChannelID int64 `json:"channel_id"` + Message string `json:"message"` + Entities []domain.MessageEntity `json:"entities"` + Media *domain.MessageMedia `json:"media"` + Silent bool `json:"silent"` + NoForwards bool `json:"noforwards"` + ReplyTo *domain.MessageReply `json:"reply_to,omitempty"` + Forward *domain.MessageForward `json:"forward,omitempty"` + ViaBotID int64 `json:"via_bot_id"` + GroupedID int64 `json:"grouped_id"` + ReplyMarkup *domain.MessageReplyMarkup `json:"reply_markup"` + RichMessage *domain.MessageRichMessage `json:"rich_message"` + SendAs *domain.Peer `json:"send_as,omitempty"` + Action *domain.ChannelMessageAction `json:"action,omitempty"` + TTLPeriod int `json:"ttl_period"` +} + +type monoforumSendFingerprintPayload struct { + Version int `json:"version"` + ChannelID int64 `json:"channel_id"` + SavedPeer domain.Peer `json:"saved_peer"` + Message string `json:"message"` + Entities []domain.MessageEntity `json:"entities"` +} + +// PrivateSendFingerprint returns a SHA-256 fingerprint of the original send +// intent. RPC callers should supply their precomputed raw-TL fingerprint; +// internal callers get a deterministic domain-level fallback. +func PrivateSendFingerprint(req domain.SendPrivateTextRequest) ([]byte, error) { + if len(req.IdempotencyFingerprint) > 0 { + if len(req.IdempotencyFingerprint) != sha256.Size { + return nil, fmt.Errorf("private send idempotency fingerprint: got %d bytes, want %d", len(req.IdempotencyFingerprint), sha256.Size) + } + return append([]byte(nil), req.IdempotencyFingerprint...), nil + } + payload, err := json.Marshal(privateSendFingerprintPayload{ + Version: privateSendFingerprintVersion, + RecipientUserID: req.RecipientUserID, + Message: req.Message, + Entities: req.Entities, + Media: req.Media, + Silent: req.Silent, + NoForwards: req.NoForwards, + ReplyTo: req.ReplyTo, + Forward: req.Forward, + TTLPeriod: req.TTLPeriod, + ViaBotID: req.ViaBotID, + GroupedID: req.GroupedID, + Effect: req.Effect, + ReplyMarkup: req.ReplyMarkup, + RichMessage: req.RichMessage, + }) + if err != nil { + return nil, fmt.Errorf("marshal private send fingerprint: %w", err) + } + sum := sha256.Sum256(payload) + return sum[:], nil +} + +// ChannelSendFingerprint returns the request-boundary SHA-256 value when present, otherwise a +// deterministic domain fallback for app/Bot API callers that do not originate from a TL request. +func ChannelSendFingerprint(req domain.SendChannelMessageRequest) ([]byte, error) { + if len(req.IdempotencyFingerprint) > 0 { + if err := ValidateSendFingerprint(req.IdempotencyFingerprint, "channel send"); err != nil { + return nil, err + } + return append([]byte(nil), req.IdempotencyFingerprint...), nil + } + payload, err := json.Marshal(channelSendFingerprintPayload{ + Version: channelSendFingerprintVersion, + ChannelID: req.ChannelID, + Message: req.Message, + Entities: req.Entities, + Media: req.Media, + Silent: req.Silent, + NoForwards: req.NoForwards, + ReplyTo: req.ReplyTo, + Forward: req.Forward, + ViaBotID: req.ViaBotID, + GroupedID: req.GroupedID, + ReplyMarkup: req.ReplyMarkup, + RichMessage: req.RichMessage, + SendAs: req.SendAs, + Action: req.Action, + TTLPeriod: req.TTLPeriod, + }) + if err != nil { + return nil, fmt.Errorf("marshal channel send fingerprint: %w", err) + } + sum := sha256.Sum256(payload) + return sum[:], nil +} + +// MonoforumSendFingerprint is scoped to one subscriber sub-dialog. SavedPeer is also part of the +// lookup key, but retaining it in the fallback prevents a future index/scope regression from +// silently accepting a cross-dialog replay. +func MonoforumSendFingerprint(req domain.SendMonoforumMessageRequest) ([]byte, error) { + if len(req.IdempotencyFingerprint) > 0 { + if err := ValidateSendFingerprint(req.IdempotencyFingerprint, "monoforum send"); err != nil { + return nil, err + } + return append([]byte(nil), req.IdempotencyFingerprint...), nil + } + payload, err := json.Marshal(monoforumSendFingerprintPayload{ + Version: channelSendFingerprintVersion, + ChannelID: req.MonoforumID, + SavedPeer: req.SavedPeer, + Message: req.Message, + Entities: req.Entities, + }) + if err != nil { + return nil, fmt.Errorf("marshal monoforum send fingerprint: %w", err) + } + sum := sha256.Sum256(payload) + return sum[:], nil +} + +// ValidateSendFingerprint rejects empty, truncated and oversized receipts. Callers must never +// guess a legacy/corrupt fingerprint from a mutable message projection. +func ValidateSendFingerprint(fingerprint []byte, operation string) error { + if len(fingerprint) != sha256.Size { + return fmt.Errorf("%s idempotency fingerprint: got %d bytes, want %d", operation, len(fingerprint), sha256.Size) + } + return nil +} + +// SameSendFingerprint requires two complete SHA-256 values. +func SameSendFingerprint(stored, expected []byte) bool { + return len(stored) == sha256.Size && len(expected) == sha256.Size && bytes.Equal(stored, expected) +} + +// SamePrivateSendFingerprint requires two complete SHA-256 values. Legacy or +// corrupt empty values are never guessed from mutable message projections. +func SamePrivateSendFingerprint(stored, expected []byte) bool { + return SameSendFingerprint(stored, expected) +} + +// EncodePrivateSendSnapshot freezes the sender-visible message returned by the +// first successful send. The snapshot is independent from mutable message-box +// projections so an edit or delete cannot erase the facts needed to acknowledge +// a later lost-response replay. +func EncodePrivateSendSnapshot(msg domain.Message) ([]byte, error) { + if msg.ID <= 0 || msg.UID <= 0 || msg.RandomID == 0 || msg.OwnerUserID == 0 || msg.Pts <= 0 { + return nil, fmt.Errorf("private send snapshot: invalid id=%d uid=%d random_id=%d owner=%d pts=%d", msg.ID, msg.UID, msg.RandomID, msg.OwnerUserID, msg.Pts) + } + raw, err := json.Marshal(privateSendSnapshotEnvelope{Version: sendSnapshotVersion, Message: msg}) + if err != nil { + return nil, fmt.Errorf("marshal private send snapshot: %w", err) + } + return raw, nil +} + +// DecodePrivateSendSnapshot returns a fresh object graph on every replay. Empty +// legacy rows are rejected rather than reconstructed from an edited projection. +func DecodePrivateSendSnapshot(raw []byte) (domain.Message, error) { + var envelope privateSendSnapshotEnvelope + if err := json.Unmarshal(raw, &envelope); err != nil { + return domain.Message{}, fmt.Errorf("unmarshal private send snapshot: %w", err) + } + msg := envelope.Message + if envelope.Version != sendSnapshotVersion || msg.ID <= 0 || msg.UID <= 0 || msg.RandomID == 0 || msg.OwnerUserID == 0 || msg.Pts <= 0 { + return domain.Message{}, fmt.Errorf("private send snapshot: invalid version=%d id=%d uid=%d random_id=%d owner=%d pts=%d", envelope.Version, msg.ID, msg.UID, msg.RandomID, msg.OwnerUserID, msg.Pts) + } + return msg, nil +} + +// EncodeChannelSendSnapshot freezes the first sender echo for random_id replay. +func EncodeChannelSendSnapshot(msg domain.ChannelMessage) ([]byte, error) { + if msg.ChannelID == 0 || msg.ID <= 0 || msg.RandomID == 0 || msg.SenderUserID == 0 || msg.Pts <= 0 { + return nil, fmt.Errorf("channel send snapshot: invalid channel=%d id=%d random_id=%d sender=%d pts=%d", msg.ChannelID, msg.ID, msg.RandomID, msg.SenderUserID, msg.Pts) + } + raw, err := json.Marshal(channelSendSnapshotEnvelope{Version: sendSnapshotVersion, Message: msg}) + if err != nil { + return nil, fmt.Errorf("marshal channel send snapshot: %w", err) + } + return raw, nil +} + +// DecodeChannelSendSnapshot returns a fresh immutable first-send projection. +func DecodeChannelSendSnapshot(raw []byte) (domain.ChannelMessage, error) { + var envelope channelSendSnapshotEnvelope + if err := json.Unmarshal(raw, &envelope); err != nil { + return domain.ChannelMessage{}, fmt.Errorf("unmarshal channel send snapshot: %w", err) + } + msg := envelope.Message + if envelope.Version != sendSnapshotVersion || msg.ChannelID == 0 || msg.ID <= 0 || msg.RandomID == 0 || msg.SenderUserID == 0 || msg.Pts <= 0 { + return domain.ChannelMessage{}, fmt.Errorf("channel send snapshot: invalid version=%d channel=%d id=%d random_id=%d sender=%d pts=%d", envelope.Version, msg.ChannelID, msg.ID, msg.RandomID, msg.SenderUserID, msg.Pts) + } + return msg, nil +} diff --git a/internal/store/private_send_idempotency_test.go b/internal/store/private_send_idempotency_test.go new file mode 100644 index 00000000..41a7de8c --- /dev/null +++ b/internal/store/private_send_idempotency_test.go @@ -0,0 +1,56 @@ +package store + +import ( + "testing" + + "telesrv/internal/domain" +) + +func TestPrivateSendSnapshotIsDeepAndVersioned(t *testing.T) { + message := domain.Message{ + ID: 7, UID: 8, RandomID: 9, OwnerUserID: 10, + Peer: domain.Peer{Type: domain.PeerTypeUser, ID: 11}, + From: domain.Peer{Type: domain.PeerTypeUser, ID: 10}, + Date: 12, Out: true, Body: "first", Pts: 13, + Entities: []domain.MessageEntity{{Type: domain.MessageEntityBold, Length: 5}}, + Media: &domain.MessageMedia{Kind: domain.MessageMediaKindContact, Contact: &domain.MessageContact{PhoneNumber: "+1", FirstName: "Alice"}}, + } + raw, err := EncodePrivateSendSnapshot(message) + if err != nil { + t.Fatalf("encode private snapshot: %v", err) + } + message.Body = "mutated" + message.Entities[0].Length = 1 + message.Media.Contact.FirstName = "Mutated" + decoded, err := DecodePrivateSendSnapshot(raw) + if err != nil { + t.Fatalf("decode private snapshot: %v", err) + } + if decoded.Body != "first" || decoded.Entities[0].Length != 5 || decoded.Media.Contact.FirstName != "Alice" { + t.Fatalf("decoded private snapshot = %+v, want immutable nested graph", decoded) + } + decoded.Media.Contact.FirstName = "Second mutation" + again, err := DecodePrivateSendSnapshot(raw) + if err != nil || again.Media.Contact.FirstName != "Alice" { + t.Fatalf("second decode = %+v err=%v, want fresh graph", again, err) + } +} + +func TestChannelSendSnapshotRejectsEmptyLegacyValue(t *testing.T) { + if _, err := DecodeChannelSendSnapshot([]byte(`{}`)); err == nil { + t.Fatal("empty legacy channel snapshot decoded successfully") + } + message := domain.ChannelMessage{ + ChannelID: 21, ID: 22, RandomID: 23, SenderUserID: 24, + From: domain.Peer{Type: domain.PeerTypeUser, ID: 24}, + Date: 25, Body: "channel", Pts: 26, + } + raw, err := EncodeChannelSendSnapshot(message) + if err != nil { + t.Fatalf("encode channel snapshot: %v", err) + } + decoded, err := DecodeChannelSendSnapshot(raw) + if err != nil || decoded.ChannelID != message.ChannelID || decoded.ID != message.ID || decoded.RandomID != message.RandomID || decoded.SenderUserID != message.SenderUserID || decoded.Body != message.Body || decoded.Pts != message.Pts { + t.Fatalf("decoded channel snapshot = %+v err=%v, want %+v", decoded, err, message) + } +} diff --git a/internal/store/redisstore/allocator_integration_test.go b/internal/store/redisstore/allocator_integration_test.go index a5a82031..84e7ff77 100644 --- a/internal/store/redisstore/allocator_integration_test.go +++ b/internal/store/redisstore/allocator_integration_test.go @@ -158,4 +158,21 @@ func TestRedisRateLimiterWindow(t *testing.T) { if allowed || retry <= 0 { t.Fatalf("AllowN second allowed=%v retry=%d, want limited with retry", allowed, retry) } + + // Heal a counter left without TTL by an older split INCRBY→EXPIRE writer. + // Without this branch one transient crash could permanently deny login-code + // issuance for the affected phone/auth-key limiter dimension. + orphanKey := key + ":orphan-no-ttl" + redisOrphanKey := rateLimitKey(orphanKey) + t.Cleanup(func() { _ = c.Del(ctx, redisOrphanKey).Err() }) + if err := c.Set(ctx, redisOrphanKey, 100, 0).Err(); err != nil { + t.Fatalf("seed no-TTL counter: %v", err) + } + allowed, retry, err = limiter.Allow(ctx, orphanKey, 1, 5*time.Second) + if err != nil || allowed || retry <= 0 || retry > 5 { + t.Fatalf("heal no-TTL counter allowed=%v retry=%d err=%v", allowed, retry, err) + } + if ttl, err := c.PTTL(ctx, redisOrphanKey).Result(); err != nil || ttl <= 0 || ttl > 5*time.Second { + t.Fatalf("healed counter TTL=%v err=%v, want (0,5s]", ttl, err) + } } diff --git a/internal/store/redisstore/code.go b/internal/store/redisstore/code.go index 3eae5454..f9616eb0 100644 --- a/internal/store/redisstore/code.go +++ b/internal/store/redisstore/code.go @@ -68,12 +68,39 @@ if not raw then redis.call('DEL', KEYS[2]) return false end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' + or tonumber(record.Version or 0) ~= tonumber(ARGV[2]) then + redis.call('DEL', KEYS[1]) + redis.call('DEL', KEYS[2]) + return false +end redis.call('DEL', KEYS[1]) redis.call('DEL', KEYS[2]) return raw ` +const updatePhoneCodeScript = ` +if redis.call('EXISTS', KEYS[1]) == 0 then + return 0 +end +redis.call('SET', KEYS[1], ARGV[1], 'KEEPTTL') +return 1 +` + +const deleteUndecodablePhoneCodeScript = ` +if redis.call('GET', KEYS[1]) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +` + func (s *CodeStore) Set(ctx context.Context, hash string, code store.PhoneCode, ttl time.Duration) error { + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return err + } + code.Revision = revision v, err := json.Marshal(code) if err != nil { return fmt.Errorf("marshal phone code: %w", err) @@ -113,25 +140,30 @@ func (s *CodeStore) Get(ctx context.Context, hash string) (store.PhoneCode, bool } var code store.PhoneCode if err := json.Unmarshal(raw, &code); err != nil { - return store.PhoneCode{}, false, fmt.Errorf("unmarshal phone code: %w", err) + // Version-zero records used JSON numbers for int64 fields. Once those + // fields became quoted strings, such a record is intentionally unusable; + // compare-and-delete it so a concurrent Set successor cannot be removed. + // Its scope cannot be decoded here; a stale scope index is harmless and is + // removed by the next scoped Set, ConsumeScoped, or VerifyScoped. + if deleteErr := s.c.Eval(ctx, deleteUndecodablePhoneCodeScript, []string{codeKey(hash)}, string(raw)).Err(); deleteErr != nil { + return store.PhoneCode{}, false, fmt.Errorf("delete undecodable phone code after %v: %w", err, deleteErr) + } + return store.PhoneCode{}, false, nil } return code, true, nil } func (s *CodeStore) Update(ctx context.Context, hash string, code store.PhoneCode) error { - key := codeKey(hash) - ttl, err := s.c.PTTL(ctx, key).Result() + revision, err := store.NewPhoneCodeRevisionToken() if err != nil { - return fmt.Errorf("redis ttl phone code: %w", err) - } - if ttl <= 0 { - return nil + return err } + code.Revision = revision v, err := json.Marshal(code) if err != nil { return fmt.Errorf("marshal phone code: %w", err) } - if err := s.c.Set(ctx, key, v, ttl).Err(); err != nil { + if err := s.c.Eval(ctx, updatePhoneCodeScript, []string{codeKey(hash)}, string(v)).Err(); err != nil { return fmt.Errorf("redis update phone code: %w", err) } return nil @@ -169,6 +201,7 @@ func (s *CodeStore) ConsumeScoped(ctx context.Context, hash string, scope store. consumeScopedCodeScript, []string{codeKey(hash), codeScopeKey(scope)}, hash, + store.PhoneCodeVersionCurrent, ).Result() if err != nil { if errors.Is(err, redis.Nil) { @@ -187,7 +220,7 @@ func (s *CodeStore) ConsumeScoped(ctx context.Context, hash string, scope store. if err := json.Unmarshal([]byte(raw), &code); err != nil { return store.PhoneCode{}, false, fmt.Errorf("unmarshal consumed phone code: %w", err) } - if code.Scope() != scope { + if code.Version != store.PhoneCodeVersionCurrent || code.Scope() != scope { return store.PhoneCode{}, false, fmt.Errorf("consumed phone code scope mismatch") } return code, true, nil diff --git a/internal/store/redisstore/code_cas.go b/internal/store/redisstore/code_cas.go new file mode 100644 index 00000000..7613532d --- /dev/null +++ b/internal/store/redisstore/code_cas.go @@ -0,0 +1,147 @@ +package redisstore + +import ( + "context" + "encoding/json" + "errors" + "fmt" + + "github.com/redis/go-redis/v9" + + "telesrv/internal/store" +) + +const getPhoneCodeSnapshotScript = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return '' +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' + or tonumber(record.Version or 0) ~= tonumber(ARGV[1]) + or (record.Revision or '') == '' then + redis.call('DEL', KEYS[1]) + return '' +end +return raw +` + +const compareAndUpdatePhoneCodeScript = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return 0 +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' + or tonumber(record.Version or 0) ~= tonumber(ARGV[1]) + or (record.Revision or '') == '' then + redis.call('DEL', KEYS[1]) + return 0 +end +if (record.Purpose or '') ~= '' or record.Revision ~= ARGV[2] then + return 0 +end +redis.call('SET', KEYS[1], ARGV[3], 'KEEPTTL') +return 1 +` + +const compareAndDeletePhoneCodeScript = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return 0 +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' + or tonumber(record.Version or 0) ~= tonumber(ARGV[1]) + or (record.Revision or '') == '' then + redis.call('DEL', KEYS[1]) + return 0 +end +if (record.Purpose or '') ~= '' or record.Revision ~= ARGV[2] then + return 0 +end +redis.call('DEL', KEYS[1]) +return 1 +` + +func (s *CodeStore) GetSnapshot(ctx context.Context, hash string) (store.PhoneCodeSnapshot, bool, error) { + value, err := s.c.Eval( + ctx, + getPhoneCodeSnapshotScript, + []string{codeKey(hash)}, + store.PhoneCodeVersionCurrent, + ).Result() + if err != nil { + if errors.Is(err, redis.Nil) { + return store.PhoneCodeSnapshot{}, false, nil + } + return store.PhoneCodeSnapshot{}, false, fmt.Errorf("redis get phone code snapshot: %w", err) + } + raw, ok := value.(string) + if !ok { + return store.PhoneCodeSnapshot{}, false, fmt.Errorf("redis get phone code snapshot: unexpected result %T", value) + } + if raw == "" { + return store.PhoneCodeSnapshot{}, false, nil + } + var record store.PhoneCode + if err := json.Unmarshal([]byte(raw), &record); err != nil { + return store.PhoneCodeSnapshot{}, false, fmt.Errorf("redis get phone code snapshot decode: %w", err) + } + if record.Version != store.PhoneCodeVersionCurrent || record.Revision == "" { + return store.PhoneCodeSnapshot{}, false, fmt.Errorf("redis get phone code snapshot returned invalid version/revision") + } + return store.PhoneCodeSnapshot{Record: record, Revision: record.Revision}, true, nil +} + +func (s *CodeStore) CompareAndUpdate(ctx context.Context, hash, expectedRevision string, next store.PhoneCode) (bool, error) { + if expectedRevision == "" || next.Version != store.PhoneCodeVersionCurrent || next.Purpose != "" { + return false, nil + } + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return false, err + } + next.Revision = revision + raw, err := json.Marshal(next) + if err != nil { + return false, fmt.Errorf("marshal compare-and-update phone code: %w", err) + } + value, err := s.c.Eval( + ctx, + compareAndUpdatePhoneCodeScript, + []string{codeKey(hash)}, + store.PhoneCodeVersionCurrent, + expectedRevision, + string(raw), + ).Result() + if err != nil { + return false, fmt.Errorf("redis compare-and-update phone code: %w", err) + } + return redisCASApplied(value, "compare-and-update phone code") +} + +func (s *CodeStore) CompareAndDelete(ctx context.Context, hash, expectedRevision string) (bool, error) { + if expectedRevision == "" { + return false, nil + } + value, err := s.c.Eval( + ctx, + compareAndDeletePhoneCodeScript, + []string{codeKey(hash)}, + store.PhoneCodeVersionCurrent, + expectedRevision, + ).Result() + if err != nil { + return false, fmt.Errorf("redis compare-and-delete phone code: %w", err) + } + return redisCASApplied(value, "compare-and-delete phone code") +} + +func redisCASApplied(value any, operation string) (bool, error) { + number, ok := value.(int64) + if !ok || (number != 0 && number != 1) { + return false, fmt.Errorf("redis %s: unexpected result %v (%T)", operation, value, value) + } + return number == 1, nil +} diff --git a/internal/store/redisstore/code_cas_integration_test.go b/internal/store/redisstore/code_cas_integration_test.go new file mode 100644 index 00000000..2f9ed5a0 --- /dev/null +++ b/internal/store/redisstore/code_cas_integration_test.go @@ -0,0 +1,237 @@ +package redisstore + +import ( + "context" + "encoding/json" + "fmt" + "sync" + "testing" + "time" + + "telesrv/internal/store" +) + +func TestRedisCodeStoreRevisionCAS(t *testing.T) { + codes, client, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + record := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016301", + Code: "111111", + Channel: "email_setup", + PendingEmail: "first@example.test", + MaxAttempts: 5, + } + key := hash("email-fixed") + if err := codes.Set(ctx, key, record, 45*time.Second); err != nil { + t.Fatal(err) + } + snapshot, found, err := codes.GetSnapshot(ctx, key) + if err != nil || !found || snapshot.Revision == "" || snapshot.Record.Revision != snapshot.Revision { + t.Fatalf("snapshot=%+v found=%v err=%v", snapshot, found, err) + } + before, err := client.PTTL(ctx, codeKey(key)).Result() + if err != nil { + t.Fatal(err) + } + next := snapshot.Record + next.Code = "222222" + next.Attempts = 1 + if applied, err := codes.CompareAndUpdate(ctx, key, "stale-token", next); err != nil || applied { + t.Fatalf("wrong-token update applied=%v err=%v", applied, err) + } + if applied, err := codes.CompareAndUpdate(ctx, key, snapshot.Revision, next); err != nil || !applied { + t.Fatalf("current update applied=%v err=%v", applied, err) + } + updated, found, err := codes.GetSnapshot(ctx, key) + if err != nil || !found || updated.Record.Code != next.Code || updated.Record.Attempts != 1 || updated.Revision == snapshot.Revision { + t.Fatalf("updated=%+v found=%v err=%v", updated, found, err) + } + after, err := client.PTTL(ctx, codeKey(key)).Result() + if err != nil || after <= 0 || after > before || before-after > 2*time.Second { + t.Fatalf("CAS TTL before=%v after=%v err=%v", before, after, err) + } + if applied, err := codes.CompareAndDelete(ctx, key, snapshot.Revision); err != nil || applied { + t.Fatalf("stale delete applied=%v err=%v", applied, err) + } + if applied, err := codes.CompareAndDelete(ctx, key, updated.Revision); err != nil || !applied { + t.Fatalf("current delete applied=%v err=%v", applied, err) + } + assertRedisCodeMissing(t, ctx, codes, key) +} + +func TestRedisCodeStoreRevisionCASFailClosedAndScopeIsolation(t *testing.T) { + codes, client, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + legacyHash := hash("legacy") + legacy := store.PhoneCode{ + Version: 0, + Revision: "legacy-revision", + Phone: "15550016302", + Code: "12345", + } + raw, err := json.Marshal(legacy) + if err != nil { + t.Fatal(err) + } + if err := client.Set(ctx, codeKey(legacyHash), raw, time.Minute).Err(); err != nil { + t.Fatal(err) + } + if _, found, err := codes.GetSnapshot(ctx, legacyHash); err != nil || found { + t.Fatalf("legacy snapshot found=%v err=%v", found, err) + } + assertRedisCodeMissing(t, ctx, codes, legacyHash) + + noRevisionHash := hash("no-revision") + legacy.Version = store.PhoneCodeVersionCurrent + legacy.Revision = "" + raw, err = json.Marshal(legacy) + if err != nil { + t.Fatal(err) + } + if err := client.Set(ctx, codeKey(noRevisionHash), raw, time.Minute).Err(); err != nil { + t.Fatal(err) + } + if _, found, err := codes.GetSnapshot(ctx, noRevisionHash); err != nil || found { + t.Fatalf("revisionless snapshot found=%v err=%v", found, err) + } + assertRedisCodeMissing(t, ctx, codes, noRevisionHash) + + scopedHash := hash("scoped") + scoped := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016303", + Code: "12345", + Purpose: store.PhoneCodePurposeChangePhone, + UserID: 42, + AuthKeyID: [8]byte{1}, + } + if err := codes.Set(ctx, scopedHash, scoped, time.Minute); err != nil { + t.Fatal(err) + } + snapshot, found, err := codes.GetSnapshot(ctx, scopedHash) + if err != nil || !found { + t.Fatalf("scoped snapshot found=%v err=%v", found, err) + } + if applied, err := codes.CompareAndUpdate(ctx, scopedHash, snapshot.Revision, snapshot.Record); err != nil || applied { + t.Fatalf("scoped update applied=%v err=%v", applied, err) + } + if applied, err := codes.CompareAndDelete(ctx, scopedHash, snapshot.Revision); err != nil || applied { + t.Fatalf("scoped delete applied=%v err=%v", applied, err) + } + if _, found, _ := codes.Get(ctx, scopedHash); !found { + t.Fatal("generic CAS mutated scoped record") + } +} + +func TestRedisCodeStoreRevisionCASPreventsABAAndHasSingleWinner(t *testing.T) { + codes, _, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + record := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016304", + Code: "123456", + Channel: "email_change", + } + key := hash("aba") + if err := codes.Set(ctx, key, record, time.Minute); err != nil { + t.Fatal(err) + } + old, found, err := codes.GetSnapshot(ctx, key) + if err != nil || !found { + t.Fatalf("old snapshot found=%v err=%v", found, err) + } + if err := codes.Set(ctx, key, record, time.Minute); err != nil { + t.Fatal(err) + } + current, found, err := codes.GetSnapshot(ctx, key) + if err != nil || !found || current.Revision == old.Revision { + t.Fatalf("replacement current=%+v old=%+v found=%v err=%v", current, old, found, err) + } + if applied, err := codes.CompareAndDelete(ctx, key, old.Revision); err != nil || applied { + t.Fatalf("ABA stale delete applied=%v err=%v", applied, err) + } + + const workers = 48 + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(index int) { + defer wg.Done() + next := current.Record + next.Code = fmt.Sprintf("%06d", index) + applied, err := codes.CompareAndUpdate(ctx, key, current.Revision, next) + if err != nil { + errs <- err + return + } + results <- applied + }(i) + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("concurrent CAS update: %v", err) + } + if winners := countRedisTrue(results); winners != 1 { + t.Fatalf("concurrent update winners=%d, want 1", winners) + } + winner, found, err := codes.GetSnapshot(ctx, key) + if err != nil || !found || winner.Revision == current.Revision { + t.Fatalf("winner=%+v found=%v err=%v", winner, found, err) + } + + results = make(chan bool, workers) + errs = make(chan error, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + applied, err := codes.CompareAndDelete(ctx, key, winner.Revision) + if err != nil { + errs <- err + return + } + results <- applied + }() + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("concurrent CAS delete: %v", err) + } + if winners := countRedisTrue(results); winners != 1 { + t.Fatalf("concurrent delete winners=%d, want 1", winners) + } +} + +func TestRedisCodeStoreLegacyUpdateCannotResurrectConsumedKey(t *testing.T) { + codes, _, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + record := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016305", + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + } + key := hash("no-resurrection") + if err := codes.Set(ctx, key, record, time.Minute); err != nil { + t.Fatal(err) + } + stale, found, err := codes.Get(ctx, key) + if err != nil || !found { + t.Fatalf("load stale record found=%v err=%v", found, err) + } + if _, found, err := codes.TakeLoginCode(ctx, key, record.Phone); err != nil || !found { + t.Fatalf("consume before stale update found=%v err=%v", found, err) + } + stale.Attempts++ + if err := codes.Update(ctx, key, stale); err != nil { + t.Fatalf("stale legacy Update: %v", err) + } + assertRedisCodeMissing(t, ctx, codes, key) +} diff --git a/internal/store/redisstore/code_integration_test.go b/internal/store/redisstore/code_integration_test.go index 4b970c0c..7b551787 100644 --- a/internal/store/redisstore/code_integration_test.go +++ b/internal/store/redisstore/code_integration_test.go @@ -27,6 +27,7 @@ func TestRedisCodeStoreScopedRotationAndSingleConsume(t *testing.T) { oldHash := fmt.Sprintf("scope-old-%d", suffix) newHash := fmt.Sprintf("scope-new-%d", suffix) rec := store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, Phone: fmt.Sprintf("1555%d", suffix), Code: "12345", Purpose: store.PhoneCodePurposeChangePhone, @@ -81,3 +82,34 @@ func TestRedisCodeStoreScopedRotationAndSingleConsume(t *testing.T) { t.Fatalf("remaining redis keys=%d err=%v", exists, err) } } + +func TestRedisCodeStoreConsumeScopedRejectsAndDeletesLegacyVersion(t *testing.T) { + addr := os.Getenv("TELESRV_TEST_REDIS_ADDR") + if addr == "" { + t.Skip("set TELESRV_TEST_REDIS_ADDR to run redis integration test") + } + ctx := context.Background() + c, err := Open(ctx, addr, "", 0) + if err != nil { + t.Fatalf("open: %v", err) + } + t.Cleanup(func() { _ = c.Close() }) + + hash := fmt.Sprintf("legacy-scope-%d", time.Now().UnixNano()) + rec := store.PhoneCode{ + Version: 0, Phone: "15550015004", Code: "12345", + Purpose: store.PhoneCodePurposeChangePhone, UserID: 44, AuthKeyID: [8]byte{4}, + } + scopeKey := codeScopeKey(rec.Scope()) + t.Cleanup(func() { _ = c.Del(ctx, codeKey(hash), scopeKey).Err() }) + codes := NewCodeStore(c) + if err := codes.Set(ctx, hash, rec, time.Minute); err != nil { + t.Fatal(err) + } + if _, found, err := codes.ConsumeScoped(ctx, hash, rec.Scope()); err != nil || found { + t.Fatalf("legacy scoped consume found=%v err=%v, want false/nil", found, err) + } + if exists, err := c.Exists(ctx, codeKey(hash), scopeKey).Result(); err != nil || exists != 0 { + t.Fatalf("legacy scoped keys remain=%d err=%v", exists, err) + } +} diff --git a/internal/store/redisstore/login_code.go b/internal/store/redisstore/login_code.go new file mode 100644 index 00000000..f282a609 --- /dev/null +++ b/internal/store/redisstore/login_code.go @@ -0,0 +1,385 @@ +package redisstore + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "strconv" + + "github.com/redis/go-redis/v9" + + "telesrv/internal/store" +) + +const verifyLoginCodeScript = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return {0, ''} +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' then + redis.call('DEL', KEYS[1]) + return {0, ''} +end +if tonumber(record.Version or 0) ~= tonumber(ARGV[1]) then + redis.call('DEL', KEYS[1]) + return {0, ''} +end +if record.SignUpVerified == true then + return {0, ''} +end +local channel = record.Channel or '' +if (record.Purpose or '') ~= '' or (record.Phone or '') ~= ARGV[2] + or (channel ~= ARGV[6] and channel ~= ARGV[7]) + or (record.Code or '') == '' or ARGV[3] == '' then + return {1, raw} +end +if (record.Code or '') ~= ARGV[3] then + local attempts = tonumber(record.Attempts or 0) + 1 + record.Attempts = attempts + record.Revision = ARGV[8] + local max_attempts = tonumber(record.MaxAttempts or 0) + if not max_attempts or max_attempts <= 0 then + max_attempts = tonumber(ARGV[5]) or 0 + end + if max_attempts <= 0 then + max_attempts = 1 + end + local updated = cjson.encode(record) + if attempts >= max_attempts then + redis.call('DEL', KEYS[1]) + else + redis.call('SET', KEYS[1], updated, 'KEEPTTL') + end + return {1, updated} +end +if ARGV[4] == '1' then + if tonumber(record.IssuedUserID or '0') ~= 0 then + return {1, raw} + end + record.SignUpVerified = true + record.Revision = ARGV[8] + local updated = cjson.encode(record) + redis.call('SET', KEYS[1], updated, 'KEEPTTL') + return {2, updated} +end +redis.call('DEL', KEYS[1]) +return {2, raw} +` + +const verifyScopedCodeScript = ` +if redis.call('GET', KEYS[2]) ~= ARGV[1] then + return {0, ''} +end +local raw = redis.call('GET', KEYS[1]) +if not raw then + redis.call('DEL', KEYS[2]) + return {0, ''} +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' then + redis.call('DEL', KEYS[1], KEYS[2]) + return {0, ''} +end +if tonumber(record.Version or 0) ~= tonumber(ARGV[2]) then + redis.call('DEL', KEYS[1], KEYS[2]) + return {0, ''} +end +local encoded_auth_key = '' +if type(record.AuthKeyID) == 'table' then + encoded_auth_key = cjson.encode(record.AuthKeyID) +end +if (record.Purpose or '') ~= ARGV[6] + or tonumber(record.UserID or 0) ~= tonumber(ARGV[7]) + or encoded_auth_key ~= ARGV[8] + or (record.Phone or '') ~= ARGV[9] + or record.SignUpVerified == true + or (record.Code or '') == '' then + redis.call('DEL', KEYS[1], KEYS[2]) + return {0, ''} +end +if ARGV[3] == '' then + return {1, raw} +end +if (record.Code or '') ~= ARGV[3] then + local attempts = tonumber(record.Attempts or 0) + 1 + record.Attempts = attempts + record.Revision = ARGV[5] + local max_attempts = tonumber(record.MaxAttempts or 0) + if not max_attempts or max_attempts <= 0 then + max_attempts = tonumber(ARGV[4]) or 0 + end + if max_attempts <= 0 then + max_attempts = 1 + end + local updated = cjson.encode(record) + if attempts >= max_attempts then + redis.call('DEL', KEYS[1], KEYS[2]) + else + redis.call('SET', KEYS[1], updated, 'KEEPTTL') + end + return {1, updated} +end +redis.call('DEL', KEYS[1], KEYS[2]) +return {2, raw} +` + +const takeLoginCodeScript = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return '' +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' then + redis.call('DEL', KEYS[1]) + return '' +end +if tonumber(record.Version or 0) ~= tonumber(ARGV[1]) then + redis.call('DEL', KEYS[1]) + return '' +end +if record.SignUpVerified == true then + return '' +end +local channel = record.Channel or '' +if (record.Purpose or '') ~= '' or (record.Phone or '') ~= ARGV[2] + or (channel ~= ARGV[3] and channel ~= ARGV[4] and channel ~= ARGV[5]) then + return '' +end +redis.call('DEL', KEYS[1]) +return raw +` + +const consumeSignUpVerifiedScript = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return '' +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' then + redis.call('DEL', KEYS[1]) + return '' +end +if tonumber(record.Version or 0) ~= tonumber(ARGV[1]) then + redis.call('DEL', KEYS[1]) + return '' +end +local channel = record.Channel or '' +if (record.Purpose or '') ~= '' or (record.Phone or '') ~= ARGV[2] + or (channel ~= ARGV[3] and channel ~= ARGV[4]) + or tonumber(record.IssuedUserID or '0') ~= 0 + or record.SignUpVerified ~= true then + return '' +end +redis.call('DEL', KEYS[1]) +return raw +` + +const invalidateLoginCodeScript = ` +local raw = redis.call('GET', KEYS[1]) +if not raw then + return '' +end +local decoded, record = pcall(cjson.decode, raw) +if not decoded or type(record) ~= 'table' then + redis.call('DEL', KEYS[1]) + return '' +end +if tonumber(record.Version or 0) ~= tonumber(ARGV[1]) then + redis.call('DEL', KEYS[1]) + return '' +end +local channel = record.Channel or '' +if (record.Purpose or '') ~= '' or (record.Phone or '') ~= ARGV[2] + or (channel ~= ARGV[3] and channel ~= ARGV[4] and channel ~= ARGV[5]) then + return '' +end +redis.call('DEL', KEYS[1]) +return raw +` + +func (s *CodeStore) VerifyLogin(ctx context.Context, hash, phone, code string, keepForSignUp bool, defaultMaxAttempts int) (store.LoginCodeVerifyResult, error) { + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return store.LoginCodeVerifyResult{}, err + } + keep := 0 + if keepForSignUp { + keep = 1 + } + value, err := s.c.Eval( + ctx, + verifyLoginCodeScript, + []string{codeKey(hash)}, + store.PhoneCodeVersionCurrent, + phone, + code, + keep, + defaultMaxAttempts, + store.PhoneCodeChannelPhone, + store.PhoneCodeChannelEmailLogin, + revision, + ).Result() + if err != nil { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify login code: %w", err) + } + return decodeRedisLoginCodeVerification(value) +} + +func (s *CodeStore) VerifyScoped(ctx context.Context, hash string, scope store.PhoneCodeScope, code string, defaultMaxAttempts int) (store.LoginCodeVerifyResult, error) { + if !scope.Valid() { + return store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyMissing}, nil + } + revision, err := store.NewPhoneCodeRevisionToken() + if err != nil { + return store.LoginCodeVerifyResult{}, err + } + authKeyID, err := json.Marshal(scope.AuthKeyID) + if err != nil { + return store.LoginCodeVerifyResult{}, fmt.Errorf("marshal scoped phone code auth key: %w", err) + } + value, err := s.c.Eval( + ctx, + verifyScopedCodeScript, + []string{codeKey(hash), codeScopeKey(scope)}, + hash, + store.PhoneCodeVersionCurrent, + code, + defaultMaxAttempts, + revision, + scope.Purpose, + strconv.FormatInt(scope.UserID, 10), + string(authKeyID), + scope.Phone, + ).Result() + if err != nil { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify scoped phone code: %w", err) + } + result, err := decodeRedisLoginCodeVerification(value) + if err != nil { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify scoped phone code: %w", err) + } + if result.Status != store.LoginCodeVerifyMissing && result.Record.Scope() != scope { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify scoped phone code returned a record outside the requested scope") + } + return result, nil +} + +func (s *CodeStore) ConsumeSignUpVerified(ctx context.Context, hash, phone string) (store.PhoneCode, bool, error) { + return s.consumeLoginCode( + ctx, + consumeSignUpVerifiedScript, + "consume sign-up verified code", + hash, + phone, + true, + true, + store.PhoneCodeChannelPhone, + store.PhoneCodeChannelEmailLogin, + ) +} + +func (s *CodeStore) TakeLoginCode(ctx context.Context, hash, phone string) (store.PhoneCode, bool, error) { + return s.consumeLoginCode( + ctx, + takeLoginCodeScript, + "take login code", + hash, + phone, + false, + false, + store.PhoneCodeChannelPhone, + store.PhoneCodeChannelEmailLogin, + store.PhoneCodeChannelEmailSetupRequired, + ) +} + +func (s *CodeStore) InvalidateLoginCode(ctx context.Context, hash, phone string) (bool, error) { + _, found, err := s.consumeLoginCode( + ctx, + invalidateLoginCodeScript, + "invalidate login code", + hash, + phone, + false, + true, + store.PhoneCodeChannelPhone, + store.PhoneCodeChannelEmailLogin, + store.PhoneCodeChannelEmailSetupRequired, + ) + return found, err +} + +func (s *CodeStore) consumeLoginCode(ctx context.Context, script, operation, hash, phone string, requireVerified, allowVerified bool, channels ...string) (store.PhoneCode, bool, error) { + args := make([]any, 0, 2+len(channels)) + args = append(args, store.PhoneCodeVersionCurrent, phone) + for _, channel := range channels { + args = append(args, channel) + } + value, err := s.c.Eval( + ctx, + script, + []string{codeKey(hash)}, + args..., + ).Result() + if err != nil { + if errors.Is(err, redis.Nil) { + return store.PhoneCode{}, false, nil + } + return store.PhoneCode{}, false, fmt.Errorf("redis %s: %w", operation, err) + } + raw, ok := value.(string) + if !ok { + return store.PhoneCode{}, false, fmt.Errorf("redis %s: unexpected result %T", operation, value) + } + if raw == "" { + return store.PhoneCode{}, false, nil + } + var record store.PhoneCode + if err := json.Unmarshal([]byte(raw), &record); err != nil { + return store.PhoneCode{}, false, fmt.Errorf("redis %s decode: %w", operation, err) + } + if record.Version != store.PhoneCodeVersionCurrent || record.Purpose != "" || record.Phone != phone || + !loginCodeChannelAllowed(record.Channel, channels) || + (requireVerified && (record.IssuedUserID != 0 || !record.SignUpVerified)) || + (!requireVerified && !allowVerified && record.SignUpVerified) { + return store.PhoneCode{}, false, fmt.Errorf("redis %s returned a record outside the requested login scope", operation) + } + return record, true, nil +} + +func loginCodeChannelAllowed(channel string, allowed []string) bool { + for _, item := range allowed { + if channel == item { + return true + } + } + return false +} + +func decodeRedisLoginCodeVerification(value any) (store.LoginCodeVerifyResult, error) { + items, ok := value.([]interface{}) + if !ok || len(items) != 2 { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify login code: unexpected result %T", value) + } + statusNumber, ok := items[0].(int64) + if !ok || statusNumber < int64(store.LoginCodeVerifyMissing) || statusNumber > int64(store.LoginCodeVerifyAccepted) { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify login code: invalid status %v", items[0]) + } + result := store.LoginCodeVerifyResult{Status: store.LoginCodeVerifyStatus(statusNumber)} + if result.Status == store.LoginCodeVerifyMissing { + return result, nil + } + raw, ok := items[1].(string) + if !ok || raw == "" { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify login code: status %d has invalid record %T", result.Status, items[1]) + } + if err := json.Unmarshal([]byte(raw), &result.Record); err != nil { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify login code decode: %w", err) + } + if result.Record.Version != store.PhoneCodeVersionCurrent { + return store.LoginCodeVerifyResult{}, fmt.Errorf("redis verify login code returned version %d", result.Record.Version) + } + return result, nil +} diff --git a/internal/store/redisstore/login_code_integration_test.go b/internal/store/redisstore/login_code_integration_test.go new file mode 100644 index 00000000..eff488bd --- /dev/null +++ b/internal/store/redisstore/login_code_integration_test.go @@ -0,0 +1,517 @@ +package redisstore + +import ( + "context" + "fmt" + "math" + "os" + "sync" + "testing" + "time" + + "github.com/redis/go-redis/v9" + + "telesrv/internal/store" +) + +func TestRedisCodeStoreAtomicLoginStateMachine(t *testing.T) { + codes, client, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + const phone = "15550016101" + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + IssuedUserID: 1000000101, + Phone: phone, + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + MaxAttempts: 2, + } + } + + t.Run("version and corrupt records fail closed", func(t *testing.T) { + legacy := newRecord() + legacy.Version = 0 + verifyHash := hash("legacy-verify") + if err := codes.Set(ctx, verifyHash, legacy, time.Minute); err != nil { + t.Fatal(err) + } + result, err := codes.VerifyLogin(ctx, verifyHash, phone, legacy.Code, false, 5) + if err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("legacy verify = %+v err=%v", result, err) + } + assertRedisCodeMissing(t, ctx, codes, verifyHash) + + // This is the actual pre-state-machine JSON shape: int64 fields were + // numbers, not the quoted strings emitted by current Set. + legacyRaw := fmt.Sprintf( + `{"Version":0,"IssuedUserID":1000000101,"Phone":%q,"Code":"12345","Channel":"phone","UserID":42,"SessionID":9007199254740993}`, + phone, + ) + getLegacyHash := hash("legacy-raw-number-get") + if err := client.Set(ctx, codeKey(getLegacyHash), legacyRaw, time.Minute).Err(); err != nil { + t.Fatal(err) + } + if got, found, err := codes.Get(ctx, getLegacyHash); err != nil || found { + t.Fatalf("legacy raw-number Get=%+v found=%v err=%v, want Missing", got, found, err) + } + if exists, err := client.Exists(ctx, codeKey(getLegacyHash)).Result(); err != nil || exists != 0 { + t.Fatalf("legacy raw-number Get left key exists=%d err=%v", exists, err) + } + + verifyLegacyHash := hash("legacy-raw-number-verify") + if err := client.Set(ctx, codeKey(verifyLegacyHash), legacyRaw, time.Minute).Err(); err != nil { + t.Fatal(err) + } + result, err = codes.VerifyLogin(ctx, verifyLegacyHash, phone, "12345", false, 5) + if err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("legacy raw-number VerifyLogin=%+v err=%v, want Missing", result, err) + } + if exists, err := client.Exists(ctx, codeKey(verifyLegacyHash)).Result(); err != nil || exists != 0 { + t.Fatalf("legacy raw-number VerifyLogin left key exists=%d err=%v", exists, err) + } + + unknown := newRecord() + unknown.Version++ + takeHash := hash("unknown-take") + if err := codes.Set(ctx, takeHash, unknown, time.Minute); err != nil { + t.Fatal(err) + } + if _, found, err := codes.TakeLoginCode(ctx, takeHash, phone); err != nil || found { + t.Fatalf("unknown take found=%v err=%v", found, err) + } + assertRedisCodeMissing(t, ctx, codes, takeHash) + + legacy.SignUpVerified = true + consumeHash := hash("legacy-consume") + if err := codes.Set(ctx, consumeHash, legacy, time.Minute); err != nil { + t.Fatal(err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, consumeHash, phone); err != nil || found { + t.Fatalf("legacy sign-up consume found=%v err=%v", found, err) + } + assertRedisCodeMissing(t, ctx, codes, consumeHash) + + corruptHash := hash("corrupt") + if err := client.Set(ctx, codeKey(corruptHash), `{`, time.Minute).Err(); err != nil { + t.Fatal(err) + } + result, err = codes.VerifyLogin(ctx, corruptHash, phone, "12345", false, 5) + if err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("corrupt verify = %+v err=%v", result, err) + } + assertRedisCodeMissing(t, ctx, codes, corruptHash) + }) + + t.Run("scope channel and issued-user gates", func(t *testing.T) { + record := newRecord() + scopeHash := hash("scope") + if err := codes.Set(ctx, scopeHash, record, time.Minute); err != nil { + t.Fatal(err) + } + result, err := codes.VerifyLogin(ctx, scopeHash, "15550016999", record.Code, false, 5) + if err != nil || result.Status != store.LoginCodeVerifyInvalid || result.Record.Attempts != 0 { + t.Fatalf("cross-phone verify = %+v err=%v", result, err) + } + stored, found, err := codes.Get(ctx, scopeHash) + if err != nil || !found || stored.Attempts != 0 { + t.Fatalf("cross-phone stored=%+v found=%v err=%v", stored, found, err) + } + + wrongChannel := record + wrongChannel.Channel = "email_setup" + channelHash := hash("channel") + if err := codes.Set(ctx, channelHash, wrongChannel, time.Minute); err != nil { + t.Fatal(err) + } + result, err = codes.VerifyLogin(ctx, channelHash, phone, wrongChannel.Code, false, 5) + if err != nil || result.Status != store.LoginCodeVerifyInvalid { + t.Fatalf("wrong-channel verify = %+v err=%v", result, err) + } + if _, found, err := codes.TakeLoginCode(ctx, channelHash, phone); err != nil || found { + t.Fatalf("wrong-channel take found=%v err=%v", found, err) + } + + issuedHash := hash("issued-existing") + record.IssuedUserID = math.MaxInt64 - 1 + if err := codes.Set(ctx, issuedHash, record, time.Minute); err != nil { + t.Fatal(err) + } + result, err = codes.VerifyLogin(ctx, issuedHash, phone, record.Code, true, 5) + if err != nil || result.Status != store.LoginCodeVerifyInvalid || result.Record.IssuedUserID != record.IssuedUserID { + t.Fatalf("issued-existing keep = %+v err=%v, want exact int64 Invalid", result, err) + } + taken, found, err := codes.TakeLoginCode(ctx, issuedHash, phone) + if err != nil || !found || taken.IssuedUserID != record.IssuedUserID { + t.Fatalf("issued-existing take=%+v found=%v err=%v", taken, found, err) + } + + setupRequired := newRecord() + setupRequired.Channel = store.PhoneCodeChannelEmailSetupRequired + setupRequired.Code = "" + setupHash := hash("setup-required") + if err := codes.Set(ctx, setupHash, setupRequired, time.Minute); err != nil { + t.Fatal(err) + } + if _, found, err := codes.TakeLoginCode(ctx, setupHash, phone); err != nil || !found { + t.Fatalf("setup-required take found=%v err=%v, want true", found, err) + } + }) + + t.Run("wrong attempts and ttl are atomic", func(t *testing.T) { + record := newRecord() + wrongHash := hash("wrong") + if err := codes.Set(ctx, wrongHash, record, 45*time.Second); err != nil { + t.Fatal(err) + } + before, err := client.PTTL(ctx, codeKey(wrongHash)).Result() + if err != nil { + t.Fatal(err) + } + first, err := codes.VerifyLogin(ctx, wrongHash, phone, "00000", false, 9) + if err != nil || first.Status != store.LoginCodeVerifyInvalid || first.Record.Attempts != 1 { + t.Fatalf("first wrong = %+v err=%v", first, err) + } + after, err := client.PTTL(ctx, codeKey(wrongHash)).Result() + if err != nil || after <= 0 || after > before || before-after > 2*time.Second { + t.Fatalf("wrong-code TTL before=%v after=%v err=%v, want KEEPTTL", before, after, err) + } + second, err := codes.VerifyLogin(ctx, wrongHash, phone, "00000", false, 9) + if err != nil || second.Status != store.LoginCodeVerifyInvalid || second.Record.Attempts != 2 { + t.Fatalf("threshold wrong = %+v err=%v", second, err) + } + assertRedisCodeMissing(t, ctx, codes, wrongHash) + }) + + t.Run("consume and sign-up marker terminal states", func(t *testing.T) { + record := newRecord() + consumeHash := hash("verify-consume") + if err := codes.Set(ctx, consumeHash, record, time.Minute); err != nil { + t.Fatal(err) + } + accepted, err := codes.VerifyLogin(ctx, consumeHash, phone, record.Code, false, 5) + if err != nil || accepted.Status != store.LoginCodeVerifyAccepted { + t.Fatalf("consume verify = %+v err=%v", accepted, err) + } + assertRedisCodeMissing(t, ctx, codes, consumeHash) + + signUp := record + signUp.IssuedUserID = 0 + signUpHash := hash("signup") + if err := codes.Set(ctx, signUpHash, signUp, 45*time.Second); err != nil { + t.Fatal(err) + } + before, err := client.PTTL(ctx, codeKey(signUpHash)).Result() + if err != nil { + t.Fatal(err) + } + marked, err := codes.VerifyLogin(ctx, signUpHash, phone, signUp.Code, true, 5) + if err != nil || marked.Status != store.LoginCodeVerifyAccepted || !marked.Record.SignUpVerified { + t.Fatalf("mark sign-up = %+v err=%v", marked, err) + } + after, err := client.PTTL(ctx, codeKey(signUpHash)).Result() + if err != nil || after <= 0 || after > before || before-after > 2*time.Second { + t.Fatalf("marker TTL before=%v after=%v err=%v, want KEEPTTL", before, after, err) + } + repeated, err := codes.VerifyLogin(ctx, signUpHash, phone, signUp.Code, true, 5) + if err != nil || repeated.Status != store.LoginCodeVerifyMissing { + t.Fatalf("repeated marker verify = %+v err=%v", repeated, err) + } + if _, found, err := codes.TakeLoginCode(ctx, signUpHash, phone); err != nil || found { + t.Fatalf("terminal marker take found=%v err=%v, want false", found, err) + } + consumed, found, err := codes.ConsumeSignUpVerified(ctx, signUpHash, phone) + if err != nil || !found || !consumed.SignUpVerified || consumed.IssuedUserID != 0 { + t.Fatalf("consume marker=%+v found=%v err=%v", consumed, found, err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, signUpHash, phone); err != nil || found { + t.Fatalf("second marker consume found=%v err=%v", found, err) + } + }) +} + +func TestRedisCodeStoreAtomicLoginConcurrency(t *testing.T) { + codes, _, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + const ( + phone = "15550016102" + workers = 48 + ) + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: phone, + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + MaxAttempts: 7, + } + } + + t.Run("consume verify one accepted", func(t *testing.T) { + key := hash("verify-race") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentRedisVerify(t, codes, key, phone, "12345", false, workers) + if statuses[store.LoginCodeVerifyAccepted] != 1 || statuses[store.LoginCodeVerifyMissing] != workers-1 || statuses[store.LoginCodeVerifyInvalid] != 0 { + t.Fatalf("verify race statuses=%+v", statuses) + } + }) + + t.Run("mark and consume each one winner", func(t *testing.T) { + key := hash("signup-race") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentRedisVerify(t, codes, key, phone, "12345", true, workers) + if statuses[store.LoginCodeVerifyAccepted] != 1 || statuses[store.LoginCodeVerifyMissing] != workers-1 { + t.Fatalf("sign-up verify race statuses=%+v", statuses) + } + if found := concurrentRedisConsume(t, codes, key, phone, workers); found != 1 { + t.Fatalf("sign-up consume winners=%d, want 1", found) + } + }) + + t.Run("take one winner", func(t *testing.T) { + key := hash("take-race") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + if found := concurrentRedisTake(t, codes, key, phone, workers); found != 1 { + t.Fatalf("take winners=%d, want 1", found) + } + }) + + t.Run("wrong attempts cannot be lost", func(t *testing.T) { + key := hash("wrong-race") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentRedisVerify(t, codes, key, phone, "00000", false, workers) + if statuses[store.LoginCodeVerifyInvalid] != 7 || statuses[store.LoginCodeVerifyMissing] != workers-7 { + t.Fatalf("wrong race statuses=%+v", statuses) + } + assertRedisCodeMissing(t, ctx, codes, key) + }) + + t.Run("verify and take share one winner", func(t *testing.T) { + key := hash("mixed-race") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(take bool) { + defer wg.Done() + if take { + _, found, err := codes.TakeLoginCode(ctx, key, phone) + if err != nil { + errs <- err + return + } + results <- found + return + } + result, err := codes.VerifyLogin(ctx, key, phone, "12345", false, 5) + if err != nil { + errs <- err + return + } + results <- result.Status == store.LoginCodeVerifyAccepted + }(i%2 == 0) + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("mixed race: %v", err) + } + winners := 0 + for won := range results { + if won { + winners++ + } + } + if winners != 1 { + t.Fatalf("mixed race winners=%d, want 1", winners) + } + }) + + t.Run("sign-up mark and take share one winner", func(t *testing.T) { + key := hash("mixed-signup-race") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(take bool) { + defer wg.Done() + if take { + _, found, err := codes.TakeLoginCode(ctx, key, phone) + if err != nil { + errs <- err + return + } + results <- found + return + } + result, err := codes.VerifyLogin(ctx, key, phone, "12345", true, 5) + if err != nil { + errs <- err + return + } + results <- result.Status == store.LoginCodeVerifyAccepted + }(i%2 == 0) + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("mixed sign-up race: %v", err) + } + if winners := countRedisTrue(results); winners != 1 { + t.Fatalf("mixed sign-up race winners=%d, want 1", winners) + } + }) +} + +func newRedisLoginCodeHarness(t *testing.T) (*CodeStore, *redis.Client, func(string) string) { + t.Helper() + addr := os.Getenv("TELESRV_TEST_REDIS_ADDR") + if addr == "" { + t.Skip("set TELESRV_TEST_REDIS_ADDR to run redis login-code integration tests") + } + ctx := context.Background() + client, err := Open(ctx, addr, "", 0) + if err != nil { + t.Fatalf("open redis: %v", err) + } + prefix := fmt.Sprintf("atomic-login-%d-", time.Now().UnixNano()) + var mu sync.Mutex + keys := make([]string, 0) + newHash := func(label string) string { + value := prefix + label + mu.Lock() + keys = append(keys, codeKey(value)) + mu.Unlock() + return value + } + t.Cleanup(func() { + mu.Lock() + cleanup := append([]string(nil), keys...) + mu.Unlock() + if len(cleanup) > 0 { + _ = client.Del(context.Background(), cleanup...).Err() + } + _ = client.Close() + }) + return NewCodeStore(client), client, newHash +} + +func assertRedisCodeMissing(t *testing.T, ctx context.Context, codes *CodeStore, hash string) { + t.Helper() + if _, found, err := codes.Get(ctx, hash); err != nil || found { + t.Fatalf("code %q found=%v err=%v, want missing", hash, found, err) + } +} + +func concurrentRedisVerify(t *testing.T, codes *CodeStore, hash, phone, code string, keep bool, workers int) map[store.LoginCodeVerifyStatus]int { + t.Helper() + ctx := context.Background() + results := make(chan store.LoginCodeVerifyStatus, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + result, err := codes.VerifyLogin(ctx, hash, phone, code, keep, 5) + if err != nil { + errs <- err + return + } + results <- result.Status + }() + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("VerifyLogin: %v", err) + } + counts := make(map[store.LoginCodeVerifyStatus]int) + for status := range results { + counts[status]++ + } + return counts +} + +func concurrentRedisTake(t *testing.T, codes *CodeStore, hash, phone string, workers int) int { + t.Helper() + ctx := context.Background() + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, found, err := codes.TakeLoginCode(ctx, hash, phone) + if err != nil { + errs <- err + return + } + results <- found + }() + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("TakeLoginCode: %v", err) + } + return countRedisTrue(results) +} + +func concurrentRedisConsume(t *testing.T, codes *CodeStore, hash, phone string, workers int) int { + t.Helper() + ctx := context.Background() + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + _, found, err := codes.ConsumeSignUpVerified(ctx, hash, phone) + if err != nil { + errs <- err + return + } + results <- found + }() + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("ConsumeSignUpVerified: %v", err) + } + return countRedisTrue(results) +} + +func countRedisTrue(results <-chan bool) int { + count := 0 + for found := range results { + if found { + count++ + } + } + return count +} diff --git a/internal/store/redisstore/login_code_invalidate_integration_test.go b/internal/store/redisstore/login_code_invalidate_integration_test.go new file mode 100644 index 00000000..531a2844 --- /dev/null +++ b/internal/store/redisstore/login_code_invalidate_integration_test.go @@ -0,0 +1,103 @@ +package redisstore + +import ( + "context" + "sync" + "testing" + "time" + + "telesrv/internal/store" +) + +func TestRedisCodeStoreAtomicLoginInvalidation(t *testing.T) { + codes, _, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + const phone = "15550016111" + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: phone, + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + MaxAttempts: 5, + } + } + + t.Run("owner cleanup may delete a terminal sign-up marker", func(t *testing.T) { + key := hash("invalidate-marker") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + verified, err := codes.VerifyLogin(ctx, key, phone, "12345", true, 5) + if err != nil || verified.Status != store.LoginCodeVerifyAccepted || !verified.Record.SignUpVerified { + t.Fatalf("mark sign-up=%+v err=%v", verified, err) + } + if removed, err := codes.InvalidateLoginCode(ctx, key, "15550016999"); err != nil || removed { + t.Fatalf("cross-phone invalidate removed=%v err=%v", removed, err) + } + if removed, err := codes.InvalidateLoginCode(ctx, key, phone); err != nil || !removed { + t.Fatalf("owner invalidate removed=%v err=%v", removed, err) + } + if _, found, err := codes.ConsumeSignUpVerified(ctx, key, phone); err != nil || found { + t.Fatalf("consume after invalidate found=%v err=%v", found, err) + } + }) + + t.Run("legacy records fail closed", func(t *testing.T) { + key := hash("invalidate-legacy") + legacy := newRecord() + legacy.Version = 0 + if err := codes.Set(ctx, key, legacy, time.Minute); err != nil { + t.Fatal(err) + } + if removed, err := codes.InvalidateLoginCode(ctx, key, phone); err != nil || removed { + t.Fatalf("legacy invalidate removed=%v err=%v, want false", removed, err) + } + assertRedisCodeMissing(t, ctx, codes, key) + }) + + t.Run("invalidate and sign-up consume have one winner", func(t *testing.T) { + key := hash("invalidate-race") + if err := codes.Set(ctx, key, newRecord(), time.Minute); err != nil { + t.Fatal(err) + } + if verified, err := codes.VerifyLogin(ctx, key, phone, "12345", true, 5); err != nil || verified.Status != store.LoginCodeVerifyAccepted { + t.Fatalf("mark sign-up=%+v err=%v", verified, err) + } + + const workers = 48 + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(invalidate bool) { + defer wg.Done() + if invalidate { + removed, err := codes.InvalidateLoginCode(ctx, key, phone) + if err != nil { + errs <- err + return + } + results <- removed + return + } + _, found, err := codes.ConsumeSignUpVerified(ctx, key, phone) + if err != nil { + errs <- err + return + } + results <- found + }(i%2 == 0) + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("invalidate/consume race: %v", err) + } + if winners := countRedisTrue(results); winners != 1 { + t.Fatalf("invalidate/consume winners=%d, want 1", winners) + } + }) +} diff --git a/internal/store/redisstore/ratelimit.go b/internal/store/redisstore/ratelimit.go index 6c916ba6..407369e9 100644 --- a/internal/store/redisstore/ratelimit.go +++ b/internal/store/redisstore/ratelimit.go @@ -3,7 +3,6 @@ package redisstore import ( "context" "fmt" - "math" "time" "github.com/redis/go-redis/v9" @@ -23,6 +22,16 @@ func rateLimitKey(key string) string { return "ratelimit:" + key } +const rateLimitIncrementScript = ` +local count = redis.call('INCRBY', KEYS[1], ARGV[1]) +local ttl_ms = redis.call('PTTL', KEYS[1]) +if ttl_ms < 0 then + redis.call('PEXPIRE', KEYS[1], ARGV[2]) + ttl_ms = tonumber(ARGV[2]) +end +return {count, ttl_ms} +` + func (l *RateLimiter) Allow(ctx context.Context, key string, limit int, window time.Duration) (bool, int, error) { return l.AllowN(ctx, key, 1, limit, window) } @@ -41,24 +50,29 @@ func (l *RateLimiter) AllowN(ctx context.Context, key string, cost, limit int, w return false, 0, fmt.Errorf("redis rate limiter: nil client") } redisKey := rateLimitKey(key) - count, err := l.c.IncrBy(ctx, redisKey, int64(cost)).Result() - if err != nil { - return false, 0, fmt.Errorf("redis incrby rate limit: %w", err) + windowMillis := window.Milliseconds() + if windowMillis <= 0 { + windowMillis = 1 } - if count == int64(cost) { - if err := l.c.Expire(ctx, redisKey, window).Err(); err != nil { - return false, 0, fmt.Errorf("redis expire rate limit: %w", err) - } + value, err := l.c.Eval(ctx, rateLimitIncrementScript, []string{redisKey}, cost, windowMillis).Result() + if err != nil { + return false, 0, fmt.Errorf("redis increment rate limit: %w", err) + } + items, ok := value.([]interface{}) + if !ok || len(items) != 2 { + return false, 0, fmt.Errorf("redis increment rate limit: unexpected result %T", value) + } + count, countOK := items[0].(int64) + ttlMillis, ttlOK := items[1].(int64) + if !countOK || !ttlOK || ttlMillis <= 0 { + return false, 0, fmt.Errorf("redis increment rate limit: invalid result %#v", items) } if count <= int64(limit) { return true, 0, nil } - ttl, err := l.c.TTL(ctx, redisKey).Result() - if err != nil { - return false, 0, fmt.Errorf("redis ttl rate limit: %w", err) + retry := (ttlMillis + 999) / 1000 + if retry <= 0 { + retry = 1 } - if ttl <= 0 { - ttl = window - } - return false, int(math.Ceil(ttl.Seconds())), nil + return false, int(retry), nil } diff --git a/internal/store/redisstore/scoped_code_state_integration_test.go b/internal/store/redisstore/scoped_code_state_integration_test.go new file mode 100644 index 00000000..0f022756 --- /dev/null +++ b/internal/store/redisstore/scoped_code_state_integration_test.go @@ -0,0 +1,288 @@ +package redisstore + +import ( + "context" + "encoding/json" + "math" + "sync" + "testing" + "time" + + "github.com/redis/go-redis/v9" + + "telesrv/internal/store" +) + +func TestRedisCodeStoreAtomicScopedVerification(t *testing.T) { + codes, client, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + newRecord := func() store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016121", + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + Purpose: store.PhoneCodePurposeChangePhone, + UserID: math.MaxInt64 - 121, + AuthKeyID: [8]byte{1, 2, 3, 4}, + SessionID: math.MaxInt64 - 21, + MaxAttempts: 2, + } + } + recordForCleanup := newRecord() + t.Cleanup(func() { _ = client.Del(context.Background(), codeScopeKey(recordForCleanup.Scope())).Err() }) + + t.Run("only the active hash and exact scope can mutate", func(t *testing.T) { + record := newRecord() + oldHash := hash("scoped-old") + currentHash := hash("scoped-current") + if err := codes.Set(ctx, oldHash, record, time.Minute); err != nil { + t.Fatal(err) + } + if err := codes.Set(ctx, currentHash, record, time.Minute); err != nil { + t.Fatal(err) + } + if result, err := codes.VerifyScoped(ctx, oldHash, record.Scope(), record.Code, 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("old-hash verify=%+v err=%v", result, err) + } + + otherScope := record.Scope() + otherScope.AuthKeyID = [8]byte{9} + if result, err := codes.VerifyScoped(ctx, currentHash, otherScope, "00000", 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("cross-scope verify=%+v err=%v", result, err) + } + stored, found, err := codes.Get(ctx, currentHash) + if err != nil || !found || stored.Attempts != 0 || stored.UserID != record.UserID { + t.Fatalf("victim after cross-scope verify=%+v found=%v err=%v", stored, found, err) + } + }) + + t.Run("wrong attempts preserve ttl then delete code and index", func(t *testing.T) { + record := newRecord() + key := hash("scoped-wrong") + if err := codes.Set(ctx, key, record, 45*time.Second); err != nil { + t.Fatal(err) + } + beforeTTL, err := client.PTTL(ctx, codeKey(key)).Result() + if err != nil { + t.Fatal(err) + } + before, found, err := codes.Get(ctx, key) + if err != nil || !found { + t.Fatalf("get before found=%v err=%v", found, err) + } + first, err := codes.VerifyScoped(ctx, key, record.Scope(), "00000", 9) + if err != nil || first.Status != store.LoginCodeVerifyInvalid || first.Record.Attempts != 1 || first.Record.UserID != record.UserID || first.Record.SessionID != record.SessionID { + t.Fatalf("first wrong=%+v err=%v", first, err) + } + afterTTL, err := client.PTTL(ctx, codeKey(key)).Result() + if err != nil || afterTTL <= 0 || afterTTL > beforeTTL || beforeTTL-afterTTL > 2*time.Second { + t.Fatalf("wrong-attempt TTL before=%v after=%v err=%v", beforeTTL, afterTTL, err) + } + after, found, err := codes.Get(ctx, key) + if err != nil || !found || after.Revision == before.Revision || after.UserID != record.UserID || after.SessionID != record.SessionID { + t.Fatalf("get after=%+v found=%v err=%v", after, found, err) + } + second, err := codes.VerifyScoped(ctx, key, record.Scope(), "00000", 9) + if err != nil || second.Status != store.LoginCodeVerifyInvalid || second.Record.Attempts != 2 { + t.Fatalf("threshold wrong=%+v err=%v", second, err) + } + assertRedisScopedMissing(t, ctx, client, key, record.Scope()) + if result, err := codes.VerifyScoped(ctx, key, record.Scope(), record.Code, 9); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("verify after exhaustion=%+v err=%v", result, err) + } + }) + + t.Run("correct code consumes both keys exactly once", func(t *testing.T) { + record := newRecord() + key := hash("scoped-correct") + if err := codes.Set(ctx, key, record, time.Minute); err != nil { + t.Fatal(err) + } + expected, found, err := codes.Get(ctx, key) + if err != nil || !found { + t.Fatalf("get expected found=%v err=%v", found, err) + } + accepted, err := codes.VerifyScoped(ctx, key, record.Scope(), record.Code, 5) + if err != nil || accepted.Status != store.LoginCodeVerifyAccepted || accepted.Record != expected || accepted.Record.UserID != record.UserID { + t.Fatalf("accepted=%+v err=%v, want %+v", accepted, err, expected) + } + assertRedisScopedMissing(t, ctx, client, key, record.Scope()) + if repeated, err := codes.VerifyScoped(ctx, key, record.Scope(), record.Code, 5); err != nil || repeated.Status != store.LoginCodeVerifyMissing { + t.Fatalf("repeated verify=%+v err=%v", repeated, err) + } + }) + + t.Run("legacy corrupt and inconsistent records fail closed", func(t *testing.T) { + legacy := newRecord() + legacy.Version = 0 + legacyHash := hash("scoped-legacy") + if err := codes.Set(ctx, legacyHash, legacy, time.Minute); err != nil { + t.Fatal(err) + } + if result, err := codes.VerifyScoped(ctx, legacyHash, legacy.Scope(), legacy.Code, 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("legacy verify=%+v err=%v", result, err) + } + assertRedisScopedMissing(t, ctx, client, legacyHash, legacy.Scope()) + + corrupt := newRecord() + corruptHash := hash("scoped-corrupt") + if err := codes.Set(ctx, corruptHash, corrupt, time.Minute); err != nil { + t.Fatal(err) + } + if err := client.Set(ctx, codeKey(corruptHash), `{`, time.Minute).Err(); err != nil { + t.Fatal(err) + } + if result, err := codes.VerifyScoped(ctx, corruptHash, corrupt.Scope(), corrupt.Code, 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("corrupt verify=%+v err=%v", result, err) + } + assertRedisScopedMissing(t, ctx, client, corruptHash, corrupt.Scope()) + + inconsistent := newRecord() + inconsistentHash := hash("scoped-inconsistent") + if err := codes.Set(ctx, inconsistentHash, inconsistent, time.Minute); err != nil { + t.Fatal(err) + } + stored, found, err := codes.Get(ctx, inconsistentHash) + if err != nil || !found { + t.Fatalf("get inconsistent seed found=%v err=%v", found, err) + } + stored.Phone = "15550016999" + raw, err := json.Marshal(stored) + if err != nil { + t.Fatal(err) + } + if err := client.Set(ctx, codeKey(inconsistentHash), raw, time.Minute).Err(); err != nil { + t.Fatal(err) + } + if result, err := codes.VerifyScoped(ctx, inconsistentHash, inconsistent.Scope(), inconsistent.Code, 5); err != nil || result.Status != store.LoginCodeVerifyMissing { + t.Fatalf("inconsistent verify=%+v err=%v", result, err) + } + assertRedisScopedMissing(t, ctx, client, inconsistentHash, inconsistent.Scope()) + }) +} + +func TestRedisCodeStoreAtomicScopedConcurrency(t *testing.T) { + codes, client, hash := newRedisLoginCodeHarness(t) + ctx := context.Background() + const workers = 48 + newRecord := func(maxAttempts int) store.PhoneCode { + return store.PhoneCode{ + Version: store.PhoneCodeVersionCurrent, + Phone: "15550016122", + Code: "12345", + Channel: store.PhoneCodeChannelPhone, + Purpose: store.PhoneCodePurposeChangePhone, + UserID: math.MaxInt64 - 122, + AuthKeyID: [8]byte{5, 6, 7, 8}, + SessionID: math.MaxInt64 - 22, + MaxAttempts: maxAttempts, + } + } + cleanupRecord := newRecord(7) + t.Cleanup(func() { _ = client.Del(context.Background(), codeScopeKey(cleanupRecord.Scope())).Err() }) + + t.Run("correct verification has one winner", func(t *testing.T) { + record := newRecord(7) + key := hash("scoped-verify-race") + if err := codes.Set(ctx, key, record, time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentRedisScopedVerify(t, codes, key, record.Scope(), record.Code, workers) + if statuses[store.LoginCodeVerifyAccepted] != 1 || statuses[store.LoginCodeVerifyMissing] != workers-1 || statuses[store.LoginCodeVerifyInvalid] != 0 { + t.Fatalf("correct race statuses=%+v", statuses) + } + }) + + t.Run("wrong attempts cannot be lost", func(t *testing.T) { + record := newRecord(7) + key := hash("scoped-wrong-race") + if err := codes.Set(ctx, key, record, time.Minute); err != nil { + t.Fatal(err) + } + statuses := concurrentRedisScopedVerify(t, codes, key, record.Scope(), "00000", workers) + if statuses[store.LoginCodeVerifyInvalid] != 7 || statuses[store.LoginCodeVerifyMissing] != workers-7 { + t.Fatalf("wrong race statuses=%+v", statuses) + } + assertRedisScopedMissing(t, ctx, client, key, record.Scope()) + }) + + t.Run("verification and cancellation share one winner", func(t *testing.T) { + record := newRecord(7) + key := hash("scoped-mixed-race") + if err := codes.Set(ctx, key, record, time.Minute); err != nil { + t.Fatal(err) + } + results := make(chan bool, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func(cancel bool) { + defer wg.Done() + if cancel { + _, found, err := codes.ConsumeScoped(ctx, key, record.Scope()) + if err != nil { + errs <- err + return + } + results <- found + return + } + result, err := codes.VerifyScoped(ctx, key, record.Scope(), record.Code, 5) + if err != nil { + errs <- err + return + } + results <- result.Status == store.LoginCodeVerifyAccepted + }(i%2 == 0) + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("verify/cancel race: %v", err) + } + if winners := countRedisTrue(results); winners != 1 { + t.Fatalf("verify/cancel winners=%d, want 1", winners) + } + }) +} + +func concurrentRedisScopedVerify(t *testing.T, codes *CodeStore, hash string, scope store.PhoneCodeScope, code string, workers int) map[store.LoginCodeVerifyStatus]int { + t.Helper() + results := make(chan store.LoginCodeVerifyStatus, workers) + errs := make(chan error, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + result, err := codes.VerifyScoped(context.Background(), hash, scope, code, 5) + if err != nil { + errs <- err + return + } + results <- result.Status + }() + } + wg.Wait() + close(results) + close(errs) + for err := range errs { + t.Fatalf("VerifyScoped: %v", err) + } + statuses := make(map[store.LoginCodeVerifyStatus]int) + for status := range results { + statuses[status]++ + } + return statuses +} + +func assertRedisScopedMissing(t *testing.T, ctx context.Context, client *redis.Client, hash string, scope store.PhoneCodeScope) { + t.Helper() + exists, err := client.Exists(ctx, codeKey(hash), codeScopeKey(scope)).Result() + if err != nil || exists != 0 { + t.Fatalf("scoped keys remain=%d err=%v", exists, err) + } +} diff --git a/internal/store/send_replay.go b/internal/store/send_replay.go new file mode 100644 index 00000000..201b6730 --- /dev/null +++ b/internal/store/send_replay.go @@ -0,0 +1,20 @@ +package store + +import ( + "context" + + "telesrv/internal/domain" +) + +// PrivateSendReplayStore exposes the immutable random_id receipt independently from the send +// command. App/RPC preflight uses it before rate limits, permission gates and media/source +// resolution; SendPrivateText still owns the transactional race fence. +type PrivateSendReplayStore interface { + LookupPrivateSendReplay(ctx context.Context, req domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) +} + +// ChannelSendReplayStore is the channel/monoforum counterpart. SavedPeer in the request is part +// of the monoforum idempotency scope and is zero for ordinary channel sends. +type ChannelSendReplayStore interface { + LookupChannelSendReplay(ctx context.Context, req domain.ChannelSendReplayRequest) (domain.SendChannelMessageResult, bool, error) +} diff --git a/internal/store/storetest/groupcall.go b/internal/store/storetest/groupcall.go index 876eddff..7d688f19 100644 --- a/internal/store/storetest/groupcall.go +++ b/internal/store/storetest/groupcall.go @@ -28,6 +28,13 @@ func conferenceInviteLink(slug string) string { return links.Build(links.DefaultPublicBaseURL, "call/"+slug, url.Values{"slug": []string{slug}}) } +func contractInviteMessageID(namespace int64, suffix int) int { + // group_call_invites has a global (invitee_user_id, message_id) uniqueness + // invariant because private box IDs are per-user. Keep contract fixtures in a + // bounded, namespace-derived range so filtered/reordered test runs cannot collide. + return int(namespace%1_000_000)*1000 + suffix +} + // GroupCallStoreFactory 为每个用例提供干净的 store 与不冲突的 channel id。 type GroupCallStoreFactory func(t *testing.T) (st store.GroupCallStore, channelID int64) @@ -45,10 +52,45 @@ func RunGroupCallStoreContract(t *testing.T, factory GroupCallStoreFactory) { t.Run("ResetAllParticipants", func(t *testing.T) { contractReset(t, factory) }) t.Run("JoinVideoStateLifecycle", func(t *testing.T) { contractJoinVideoState(t, factory) }) t.Run("ConferenceChainBlocks", func(t *testing.T) { contractConferenceChainBlocks(t, factory) }) + t.Run("ConferenceInviteIdempotency", func(t *testing.T) { contractConferenceInviteIdempotency(t, factory) }) t.Run("ConferenceRecipientsTerminalAccess", func(t *testing.T) { contractConferenceRecipientsTerminalAccess(t, factory) }) t.Run("ConferenceEmptyDiscards", func(t *testing.T) { contractConferenceEmptyDiscards(t, factory) }) } +func contractConferenceInviteIdempotency(t *testing.T, factory GroupCallStoreFactory) { + st, channelID := factory(t) + ctx := context.Background() + now := baseNow() + slug := fmt.Sprintf("contract-invite-idempotency-%d", channelID) + call, err := st.CreateConferenceCall(ctx, domain.GroupCall{ + ID: channelID*100 + 62, AccessHash: channelID*100 + 69, CreatorUserID: 1, + InviteSlug: slug, InviteLink: conferenceInviteLink(slug), + RandomID: channelID*100 + 62, CreatedAt: now, + }) + if err != nil { + t.Fatalf("create conference call: %v", err) + } + want := domain.GroupCallInvite{ + CallID: call.ID, InviterUserID: 1, InviteeUserID: 4, MessageID: contractInviteMessageID(channelID, 402), + Status: domain.GroupCallInvitePending, Video: true, CreatedAt: now + 1, + } + first, err := st.CreateConferenceInvite(ctx, want) + if err != nil { + t.Fatalf("create conference invite: %v", err) + } + replay, err := st.CreateConferenceInvite(ctx, want) + if err != nil { + t.Fatalf("replay conference invite: %v", err) + } + if replay != first { + t.Fatalf("replayed conference invite = %+v, want unchanged %+v", replay, first) + } + gotCall, gotInvite, found, err := st.GetGroupCallByInviteMessage(ctx, want.InviteeUserID, want.MessageID) + if err != nil || !found || gotCall.ID != call.ID || gotInvite != first { + t.Fatalf("invite lookup = call %+v invite %+v found=%v err=%v, want call %d invite %+v", gotCall, gotInvite, found, err, call.ID, first) + } +} + func newContractCall(t *testing.T, st store.GroupCallStore, channelID, id int64) domain.GroupCall { t.Helper() call, err := st.CreateGroupCall(context.Background(), domain.GroupCall{ @@ -406,13 +448,13 @@ func contractConferenceRecipientsTerminalAccess(t *testing.T, factory GroupCallS t.Fatalf("leave historical participant: %v", err) } if _, err := st.CreateConferenceInvite(ctx, domain.GroupCallInvite{ - CallID: call.ID, InviterUserID: 1, InviteeUserID: 4, MessageID: 401, + CallID: call.ID, InviterUserID: 1, InviteeUserID: 4, MessageID: contractInviteMessageID(channelID, 401), Status: domain.GroupCallInvitePending, CreatedAt: now + 4, }); err != nil { t.Fatalf("create pending invite: %v", err) } if _, err := st.CreateConferenceInvite(ctx, domain.GroupCallInvite{ - CallID: call.ID, InviterUserID: 1, InviteeUserID: 5, MessageID: 501, + CallID: call.ID, InviterUserID: 1, InviteeUserID: 5, MessageID: contractInviteMessageID(channelID, 501), Status: domain.GroupCallInviteDeclined, CreatedAt: now + 5, UpdatedAt: now + 5, }); err != nil { t.Fatalf("create declined invite: %v", err) diff --git a/internal/store/updatestate.go b/internal/store/updatestate.go index 7ed28d3e..32b820ac 100644 --- a/internal/store/updatestate.go +++ b/internal/store/updatestate.go @@ -10,6 +10,10 @@ import ( type UpdateStateStore interface { Get(ctx context.Context, authKeyID [8]byte, userID int64) (domain.UpdateState, bool, error) Save(ctx context.Context, authKeyID [8]byte, userID int64, state domain.UpdateState) error + // ObserveClientState advances only the state that the client has proved it already owns by + // carrying it in a request (or by explicitly establishing a getState snapshot baseline). + // Durable-log retention must use this watermark, never a response state merely sent by server. + ObserveClientState(ctx context.Context, authKeyID [8]byte, userID int64, state domain.UpdateState) error Delete(ctx context.Context, authKeyID [8]byte, userID int64) error DeleteAuthKey(ctx context.Context, authKeyID [8]byte) error }