fix(stock): key pending dividends by user and event to stop duplicate buttons

Repeated /stock_portfolio calls stacked a random-token pending key per
call, and applying the dividend cleaned up only the pressed one — the
rest kept live buttons until the TTL sweep. Keys are now deterministic
(pending-dividend:<userID>:<eventID>) so re-suggesting overwrites the
previous action, retires the old message's button, and applying leaves
exactly one key to delete. Callback data carries <userID>:<eventID>
instead of a token; presses are still validated against the stored
owner, chat, and message binding.
This commit is contained in:
tiennm99 committed 2026-08-05 13:59:48 +07:00
1 parent 59534a9c4b
commit bd0ad6b817
6 files changed
+151 -112

No files matched your search

+2 -2
View File
@@ -59,12 +59,12 @@ func (s *state) handleDividendCallback(ctx context.Context, b *bot.Bot, update *
}
func (s *state) resolveDividendCallback(ctx context.Context, b *bot.Bot, query *models.CallbackQuery) (PendingDividendAction, *models.Message, string, bool, error) {
token, ok := callbackToken(query.Data)
ownerID, eventID, ok := parseDividendCallback(query.Data)
if !ok || s.pending == nil {
err := answerDividendCallback(ctx, b, query.ID, "This dividend suggestion is invalid.", true)
return PendingDividendAction{}, nil, "", true, err
}
actionKey := pendingDividendKey(token)
actionKey := pendingDividendKey(ownerID, eventID)
action, _, err := s.pending.Get(ctx, actionKey)
if errors.Is(err, storage.ErrNotFound) {
err = answerDividendCallback(ctx, b, query.ID, "This dividend suggestion expired or was already used.", true)
+62 -33
View File
@@ -55,16 +55,10 @@ func newDividendFlowState(t *testing.T, events []DividendEvent) (*state, Store,
store := storage.Typed[Portfolio](collection)
pending := storage.Typed[PendingDividendAction](collection)
fake := &fakeDividendProvider{events: events}
tokenIndex := 0
s := &state{
store: store, pending: pending,
prices: &PriceClient{HTTP: priceServer.Client(), URL: priceServer.URL},
dividends: fake, nowFn: func() time.Time { return now },
newDividendToken: func() (string, error) {
token := "abcdefghijklmnopqrstu" + string(rune('v'+tokenIndex))
tokenIndex++
return token, nil
},
}
return s, store, pending, testutil.NewRecordingBot(t), now, fake
}
@@ -233,12 +227,20 @@ func TestRecordDateCreatesSeparateActionableMessages(t *testing.T) {
t.Fatal(err)
}
}
calls := rb.Sent()
if len(calls) != 6 || !strings.Contains(calls[1].Text(), "Dividend ready to apply") || !strings.Contains(calls[2].Text(), "Dividend ready to apply") || !strings.Contains(calls[4].Text(), "Dividend ready to apply") || !strings.Contains(calls[5].Text(), "Dividend ready to apply") {
t.Fatalf("messages = %#v", calls)
suggestions, buttonRemovals := 0, 0
for _, call := range rb.Sent() {
if call.Method == "sendMessage" && strings.Contains(call.Text(), "Dividend ready to apply") {
suggestions++
}
if call.Method == "editMessageReplyMarkup" {
buttonRemovals++
}
}
if suggestions != 4 || buttonRemovals != 2 {
t.Fatalf("suggestions=%d buttonRemovals=%d calls=%#v", suggestions, buttonRemovals, rb.Sent())
}
keys, err := pending.List(context.Background(), pendingDividendPrefix)
if err != nil || len(keys) != 4 {
if err != nil || len(keys) != 2 {
t.Fatalf("pending = %v, err=%v", keys, err)
}
}
@@ -268,41 +270,50 @@ func TestDividendCallbackUsesStoredEventAndCurrentHoldingOnce(t *testing.T) {
}
}
// Re-suggesting the same event replaces the pending action instead of
// stacking a second one.
keys, err := pending.List(context.Background(), pendingDividendPrefix)
if err != nil || len(keys) != 2 {
if err != nil || len(keys) != 1 || keys[0] != pendingDividendKey(7, event.ProviderID) {
t.Fatalf("pending keys = %v, err=%v", keys, err)
}
key := keys[0]
token := strings.TrimPrefix(key, pendingDividendPrefix)
action, _, _ := pending.Get(context.Background(), key)
data := "7:" + event.ProviderID
action, _, _ := pending.Get(context.Background(), keys[0])
if action.ProviderEventID != event.ProviderID || action.Symbol != "TCB" {
t.Fatalf("pending action = %+v", action)
}
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(8, 7, action.MessageID, token)); err != nil {
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(8, 7, action.MessageID, data)); err != nil {
t.Fatal(err)
}
p, _ := LoadPortfolio(context.Background(), store, 7, now.UnixMilli())
if p.VND != 100_000 {
t.Fatalf("unauthorized click changed balance: %v", p.VND)
}
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, action.MessageID, token)); err != nil {
// The superseded first suggestion message no longer applies anything.
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, action.MessageID-2, data)); err != nil {
t.Fatal(err)
}
p, _ = LoadPortfolio(context.Background(), store, 7, now.UnixMilli())
if p.VND != 100_000 {
t.Fatalf("superseded button changed balance: %v", p.VND)
}
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, action.MessageID, data)); err != nil {
t.Fatal(err)
}
p, _ = LoadPortfolio(context.Background(), store, 7, now.UnixMilli())
if p.VND != 250_000 || p.Assets["TCB"].Base != 2_850_000 || !p.Dividends["TCB"][event.ProviderID].Processed {
t.Fatalf("processed portfolio = %+v", p)
}
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, action.MessageID, token)); err != nil {
t.Fatal(err)
// Applying the dividend removes its only pending key.
keys, err = pending.List(context.Background(), pendingDividendPrefix)
if err != nil || len(keys) != 0 {
t.Fatalf("pending keys after apply = %v, err=%v", keys, err)
}
secondAction, _, _ := pending.Get(context.Background(), keys[1])
secondToken := strings.TrimPrefix(keys[1], pendingDividendPrefix)
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, secondAction.MessageID, secondToken)); err != nil {
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, action.MessageID, data)); err != nil {
t.Fatal(err)
}
p, _ = LoadPortfolio(context.Background(), store, 7, now.UnixMilli())
if p.VND != 250_000 {
t.Fatalf("second valid button credited twice: %v", p.VND)
t.Fatalf("repeated press credited twice: %v", p.VND)
}
}
@@ -378,10 +389,10 @@ func TestDividendCallbackAppliesShareEventFromStoredHistory(t *testing.T) {
OwnerUserID: 7, ChatID: 7, MessageID: 1, ProviderEventID: "2612975", Symbol: "TCB",
PositionOpenedAt: p.Assets["TCB"].OpenedAt, CreatedAt: now.UnixMilli(), ExpiresAt: now.Add(time.Hour).UnixMilli(),
}
if err := pending.Put(context.Background(), pendingDividendKey("abcdefghijklmnopqrstuv"), action); err != nil {
if err := pending.Put(context.Background(), pendingDividendKey(7, "2612975"), action); err != nil {
t.Fatal(err)
}
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, 1, "abcdefghijklmnopqrstuv")); err != nil {
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, 1, "7:2612975")); err != nil {
t.Fatal(err)
}
p, _ = LoadPortfolio(context.Background(), store, 7, now.UnixMilli())
@@ -405,11 +416,11 @@ func TestExpiredDividendActionDoesNotApply(t *testing.T) {
OwnerUserID: 7, ChatID: 7, MessageID: 1, ProviderEventID: "2612974", Symbol: "TCB",
PositionOpenedAt: p.Assets["TCB"].OpenedAt, CreatedAt: now.Add(-2 * time.Hour).UnixMilli(), ExpiresAt: now.Add(-time.Hour).UnixMilli(),
}
key := pendingDividendKey("abcdefghijklmnopqrstuv")
key := pendingDividendKey(7, "2612974")
if err := pending.Put(context.Background(), key, action); err != nil {
t.Fatal(err)
}
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, 1, "abcdefghijklmnopqrstuv")); err != nil {
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, 1, "7:2612974")); err != nil {
t.Fatal(err)
}
p, _ = LoadPortfolio(context.Background(), store, 7, now.UnixMilli())
@@ -421,7 +432,7 @@ func TestExpiredDividendActionDoesNotApply(t *testing.T) {
}
}
func TestConcurrentPortfolioRequestsCreateSeparatePendingActions(t *testing.T) {
func TestConcurrentPortfolioRequestsKeepSinglePendingAction(t *testing.T) {
now := time.Date(2026, 6, 25, 12, 0, 0, 0, saigonLocation)
s, store, pending, rb, _, _ := newDividendFlowState(t, nil)
seedDividendFlowPortfolio(t, store, now, 100)
@@ -451,8 +462,24 @@ func TestConcurrentPortfolioRequestsCreateSeparatePendingActions(t *testing.T) {
}
}
keys, err := pending.List(context.Background(), pendingDividendPrefix)
if err != nil || len(keys) != 2 || len(rb.Sent()) != 2 {
t.Fatalf("pending=%v messages=%d err=%v", keys, len(rb.Sent()), err)
if err != nil || len(keys) != 1 || keys[0] != pendingDividendKey(7, "2612974") {
t.Fatalf("pending=%v err=%v", keys, err)
}
suggestions, buttonRemovals := 0, 0
for _, call := range rb.Sent() {
switch call.Method {
case "sendMessage":
suggestions++
case "editMessageReplyMarkup":
buttonRemovals++
}
}
if suggestions != 2 || buttonRemovals != 1 {
t.Fatalf("suggestions=%d buttonRemovals=%d calls=%#v", suggestions, buttonRemovals, rb.Sent())
}
action, _, _ := pending.Get(context.Background(), keys[0])
if action.MessageID != 2 {
t.Fatalf("surviving action should bind the newest message: %+v", action)
}
}
@@ -475,11 +502,11 @@ func TestDividendCallbackRejectsPositionOpenedAfterRecordDate(t *testing.T) {
OwnerUserID: 7, ChatID: 7, MessageID: 1, ProviderEventID: "2612974", Symbol: "TCB",
PositionOpenedAt: position.OpenedAt, CreatedAt: now.UnixMilli(), ExpiresAt: now.Add(time.Hour).UnixMilli(),
}
key := pendingDividendKey("abcdefghijklmnopqrstuv")
key := pendingDividendKey(7, "2612974")
if err := pending.Put(context.Background(), key, action); err != nil {
t.Fatal(err)
}
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, 1, "abcdefghijklmnopqrstuv")); err != nil {
if err := s.handleDividendCallback(context.Background(), rb.Bot, dividendCallbackUpdate(7, 7, 1, "7:2612974")); err != nil {
t.Fatal(err)
}
p, _ = LoadPortfolio(context.Background(), store, 7, now.UnixMilli())
@@ -501,9 +528,11 @@ func TestProviderFailureKeepsPortfolioAvailable(t *testing.T) {
}
}
func dividendCallbackUpdate(userID, chatID int64, messageID int, token string) *models.Update {
// dividendCallbackUpdate simulates a button press. data is the callback-data
// suffix after the "stock_div:" prefix, i.e. "<ownerUserID>:<eventID>".
func dividendCallbackUpdate(userID, chatID int64, messageID int, data string) *models.Update {
return &models.Update{CallbackQuery: &models.CallbackQuery{
ID: "query", From: models.User{ID: userID}, Data: dividendCallbackPrefix + token,
ID: "query", From: models.User{ID: userID}, Data: dividendCallbackPrefix + data,
Message: models.MaybeInaccessibleMessage{Type: models.MaybeInaccessibleMessageTypeMessage, Message: &models.Message{
ID: messageID, Chat: models.Chat{ID: chatID, Type: models.ChatTypePrivate},
}},
@@ -2,6 +2,7 @@ package stock
import (
"context"
"errors"
"fmt"
"html"
"sort"
@@ -289,34 +290,41 @@ func (s *state) sendDividendSuggestion(ctx context.Context, b *bot.Bot, msg *mod
if !held || position.Quantity <= 0 || position.OpenedAt != expectedOpenedAt || !positionOpenedByRecordDate(position, record) {
return nil
}
if s.pending == nil {
return errors.New("stock: pending dividend store unavailable")
}
data, ok := dividendCallbackData(userID, ref.eventID)
if !ok {
return fmt.Errorf("stock: dividend event %q cannot be encoded as callback data", ref.eventID)
}
key := pendingDividendKey(userID, ref.eventID)
previous, _, previousErr := s.pending.Get(ctx, key)
action := PendingDividendAction{
OwnerUserID: userID, ChatID: msg.Chat.ID, ProviderEventID: ref.eventID, Symbol: ref.symbol,
PositionOpenedAt: position.OpenedAt, CreatedAt: now.UnixMilli(), ExpiresAt: now.Add(pendingDividendTTL).UnixMilli(),
}
token, err := s.createPendingDividend(ctx, action)
if err != nil {
return err
}
key := pendingDividendKey(token)
sent, err := b.SendMessage(ctx, &bot.SendMessageParams{
ChatID: msg.Chat.ID,
MessageThreadID: msg.MessageThreadID,
Text: dividendEventText(record.event(ref.symbol, ref.eventID), position.Quantity, true),
ParseMode: models.ParseModeHTML,
ReplyMarkup: &models.InlineKeyboardMarkup{InlineKeyboard: [][]models.InlineKeyboardButton{{{
Text: "Apply dividend", CallbackData: dividendCallbackPrefix + token,
Text: "Apply dividend", CallbackData: data,
}}}},
})
if err != nil {
_ = s.pending.Delete(ctx, key)
return err
}
action.MessageID = sent.ID
if err := s.pending.Put(ctx, key, action); err != nil {
_ = removeDividendButton(ctx, b, action.ChatID, action.MessageID)
_ = s.pending.Delete(ctx, key)
_ = removeDividendButton(ctx, b, msg.Chat.ID, sent.ID)
return fmt.Errorf("bind dividend action message: %w", err)
}
// The Put above superseded any earlier suggestion for this event; retire
// the old message's button so only the newest suggestion is pressable.
if previousErr == nil && previous.MessageID != 0 && (previous.ChatID != msg.Chat.ID || previous.MessageID != sent.ID) {
_ = removeDividendButton(ctx, b, previous.ChatID, previous.MessageID)
}
return nil
}
+7 -8
View File
@@ -21,14 +21,13 @@ import (
// prefixes/partitions). PriceClient is reused across calls; nowFn allows
// tests to inject a deterministic clock for portfolio CreatedAt.
type state struct {
store Store
pending PendingDividendStore
prices *PriceClient
dividends DividendEventProvider
events SSIStockEventProvider
locks keylock.Map
nowFn func() time.Time
newDividendToken func() (string, error)
store Store
pending PendingDividendStore
prices *PriceClient
dividends DividendEventProvider
events SSIStockEventProvider
locks keylock.Map
nowFn func() time.Time
}
func (s *state) now() time.Time {
+38 -50
View File
@@ -2,11 +2,7 @@ package stock
import (
"context"
"crypto/rand"
"encoding/base64"
"errors"
"fmt"
"regexp"
"strconv"
"strings"
"time"
@@ -17,17 +13,16 @@ const (
dividendCallbackPrefix = "stock_div:"
pendingDividendPrefix = "pending-dividend:"
pendingDividendTTL = 24 * time.Hour
dividendTokenBytes = 16
dividendTokenLength = 22
// Telegram rejects inline-button callback data longer than 64 bytes.
maxDividendCallbackBytes = 64
)
var dividendTokenPattern = regexp.MustCompile(`^[A-Za-z0-9_-]{22}$`)
type PendingDividendStore = storage.DocStore[PendingDividendAction]
// PendingDividendAction is the server-side half of an inline button. Financial
// values live in the user's dividend history; Telegram receives only an opaque
// random token.
// values live in the user's dividend history; the button's callback data
// carries only the owner user ID and provider event ID, and every press is
// validated against the stored owner, chat, and message binding.
type PendingDividendAction struct {
OwnerUserID int64 `json:"ownerUserId" bson:"ownerUserId"`
ChatID int64 `json:"chatId" bson:"chatId"`
@@ -39,54 +34,47 @@ type PendingDividendAction struct {
ExpiresAt int64 `json:"expiresAt" bson:"expiresAt"`
}
func pendingDividendKey(token string) string { return pendingDividendPrefix + token }
func generateDividendToken() (string, error) {
raw := make([]byte, dividendTokenBytes)
if _, err := rand.Read(raw); err != nil {
return "", fmt.Errorf("generate dividend action token: %w", err)
}
return base64.RawURLEncoding.EncodeToString(raw), nil
// pendingDividendKey is deterministic per (user, event): re-suggesting the
// same dividend overwrites the previous action instead of accumulating keys,
// and applying it leaves exactly one key to delete.
func pendingDividendKey(userID int64, eventID string) string {
return pendingDividendPrefix + strconv.FormatInt(userID, 10) + ":" + eventID
}
func validDividendToken(token string) bool {
return len(token) == dividendTokenLength && dividendTokenPattern.MatchString(token)
}
func callbackToken(data string) (string, bool) {
token, ok := strings.CutPrefix(data, dividendCallbackPrefix)
if !ok || !validDividendToken(token) {
// dividendCallbackData builds the inline-button payload. It fails only when
// the provider event ID is malformed or would push the payload past
// Telegram's byte limit.
func dividendCallbackData(userID int64, eventID string) (string, bool) {
if userID <= 0 || !ssiProviderIDPattern.MatchString(eventID) {
return "", false
}
return token, true
data := dividendCallbackPrefix + strconv.FormatInt(userID, 10) + ":" + eventID
if len(data) > maxDividendCallbackBytes {
return "", false
}
return data, true
}
func (s *state) createPendingDividend(ctx context.Context, action PendingDividendAction) (string, error) {
if s.pending == nil {
return "", errors.New("stock: pending dividend store unavailable")
// parseDividendCallback recovers the owner user ID and provider event ID from
// button callback data. Callback data is client-controlled, so both parts are
// validated here and the resolved action is re-checked against the presser.
func parseDividendCallback(data string) (int64, string, bool) {
if len(data) > maxDividendCallbackBytes {
return 0, "", false
}
generator := s.newDividendToken
if generator == nil {
generator = generateDividendToken
rest, ok := strings.CutPrefix(data, dividendCallbackPrefix)
if !ok {
return 0, "", false
}
for attempt := 0; attempt < 3; attempt++ {
token, err := generator()
if err != nil {
return "", err
}
if !validDividendToken(token) {
return "", errors.New("stock: invalid generated dividend token")
}
err = s.pending.PutVersioned(ctx, pendingDividendKey(token), 0, action)
if errors.Is(err, storage.ErrConflict) {
continue
}
if err != nil {
return "", fmt.Errorf("save pending dividend action: %w", err)
}
return token, nil
ownerPart, eventID, ok := strings.Cut(rest, ":")
if !ok || !ssiProviderIDPattern.MatchString(eventID) {
return 0, "", false
}
return "", errors.New("stock: could not allocate unique dividend token")
ownerID, err := strconv.ParseInt(ownerPart, 10, 64)
if err != nil || ownerID <= 0 {
return 0, "", false
}
return ownerID, eventID, true
}
func (s *state) cleanupExpiredDividends(ctx context.Context, now int64) {
+25 -10
View File
@@ -36,9 +36,10 @@ type RecordingBot struct {
Bot *bot.Bot
Server *httptest.Server
mu sync.Mutex
calls []SentCall
failures map[string]failureResponse
mu sync.Mutex
calls []SentCall
failures map[string]failureResponse
nextMessageID int
}
type failureResponse struct {
@@ -133,6 +134,11 @@ func (rb *RecordingBot) handle(w http.ResponseWriter, r *http.Request) {
rb.mu.Lock()
rb.calls = append(rb.calls, SentCall{Method: method, Form: form})
failure, shouldFail := rb.failures[method]
messageID := rb.nextMessageID
if !shouldFail && isMessageProducingMethod(method) {
rb.nextMessageID++
messageID = rb.nextMessageID
}
rb.mu.Unlock()
w.Header().Set("Content-Type", "application/json")
@@ -141,7 +147,18 @@ func (rb *RecordingBot) handle(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(failure.body))
return
}
_, _ = w.Write([]byte(okResponseFor(method)))
_, _ = w.Write([]byte(okResponseFor(method, messageID)))
}
// isMessageProducingMethod reports whether the API method returns a Message,
// so each successful call gets a distinct incrementing message_id — mirroring
// real Telegram chats, where handlers key follow-up edits on the returned ID.
func isMessageProducingMethod(method string) bool {
switch method {
case "sendMessage", "sendSticker", "sendPhoto", "sendDocument", "sendVideo", "sendAnimation":
return true
}
return false
}
// apiMethodFromPath extracts the API method from "/bot<token>/<method>".
@@ -158,13 +175,12 @@ func apiMethodFromPath(p string) string {
// okResponseFor returns a minimal `{ok:true, result:...}` payload that the
// bot library will accept for the named API method. SendMessage / SendSticker
// expect a Message; most others accept a bool.
func okResponseFor(method string) string {
switch method {
case "sendMessage", "sendSticker", "sendPhoto", "sendDocument", "sendVideo", "sendAnimation":
func okResponseFor(method string, messageID int) string {
if isMessageProducingMethod(method) {
// Minimal shape: id, date, chat. Bot library decodes via json so
// extra fields are ignored.
msg := map[string]any{
"message_id": 1,
"message_id": messageID,
"date": 0,
"chat": map[string]any{
"id": 1,
@@ -174,9 +190,8 @@ func okResponseFor(method string) string {
body := map[string]any{"ok": true, "result": msg}
out, _ := json.Marshal(body)
return string(out)
default:
return `{"ok":true,"result":true}`
}
return `{"ok":true,"result":true}`
}
// AssertSentText fails the test if no recorded sendMessage contains the