Fix login stalls and bot redemption routing

This commit is contained in:
ShukeBta
2026-06-15 01:00:05 +08:00
committed by Shuke
parent 479a06f008
commit 406b311a8e
9 changed files with 244 additions and 18 deletions
+2 -2
View File
@@ -22,8 +22,8 @@ import (
const EnvPrefix = "MEDIASTATION"
const (
defaultDatabaseMaxOpenConns = 1
defaultDatabaseMaxIdleConns = 1
defaultDatabaseMaxOpenConns = 4
defaultDatabaseMaxIdleConns = 2
)
// Config 是根配置聚合。
+6 -3
View File
@@ -75,7 +75,7 @@ func TestEnvOverride(t *testing.T) {
}
}
func TestLoadAllowsSingleConnectionDatabaseConfig(t *testing.T) {
func TestLoadAllowsExplicitSingleConnectionDatabaseConfig(t *testing.T) {
dir := t.TempDir()
wd, _ := os.Getwd()
defer func() { _ = os.Chdir(wd) }()
@@ -90,7 +90,10 @@ func TestLoadAllowsSingleConnectionDatabaseConfig(t *testing.T) {
if err != nil {
t.Fatalf("Load() error: %v", err)
}
if cfg.Database.MaxOpenConns != defaultDatabaseMaxOpenConns {
t.Fatalf("expected MaxOpenConns=1 to remain at safe default %d, got %d", defaultDatabaseMaxOpenConns, cfg.Database.MaxOpenConns)
if cfg.Database.MaxOpenConns != 1 {
t.Fatalf("expected explicit MaxOpenConns=1 to be preserved, got %d", cfg.Database.MaxOpenConns)
}
if cfg.Database.MaxIdleConns != 1 {
t.Fatalf("expected explicit MaxIdleConns=1 to be preserved, got %d", cfg.Database.MaxIdleConns)
}
}
+53
View File
@@ -296,6 +296,59 @@ func TestLoginRetriesTransientSQLiteBusy(t *testing.T) {
}
}
func TestLoginSurvivesOneReservedSQLiteConnectionWithWALPool(t *testing.T) {
ctx := context.Background()
cfg := &config.Config{}
cfg.App.DataDir = t.TempDir()
cfg.Database.DBPath = filepath.Join(cfg.App.DataDir, "reserved-connection-login.db")
cfg.Database.WALMode = true
cfg.Database.BusyTimeout = 20
cfg.Database.MaxOpenConns = 4
cfg.Database.MaxIdleConns = 2
cfg.Secrets.JWTSecret = "test-secret"
log := zap.NewNop()
db, err := database.Open(cfg, log)
if err != nil {
t.Fatal(err)
}
sqlDB, err := db.DB()
if err != nil {
t.Fatal(err)
}
defer func() { _ = sqlDB.Close() }()
if err := database.AutoMigrate(db); err != nil {
t.Fatal(err)
}
repos := repository.New(db)
permissions := NewPermissionService(log, repos)
auth := NewAuthService(cfg, log, repos, NewTokenService(cfg, log, repos), permissions)
hash, err := bcrypt.GenerateFromPassword([]byte("password"), bcrypt.MinCost)
if err != nil {
t.Fatal(err)
}
if err := repos.User.Create(ctx, &model.User{
Username: "viewer",
PasswordHash: string(hash),
Role: "user",
Tier: "free",
IsActive: true,
}); err != nil {
t.Fatal(err)
}
reserved, err := sqlDB.Conn(ctx)
if err != nil {
t.Fatal(err)
}
defer func() { _ = reserved.Close() }()
loginCtx, cancel := context.WithTimeout(ctx, time.Second)
defer cancel()
if _, err := auth.Login(loginCtx, "viewer", "password"); err != nil {
t.Fatalf("login should not be blocked by one reserved sqlite connection: %v", err)
}
}
func TestLoginReturnsTokensWhenSQLiteWriteLockPersists(t *testing.T) {
ctx := context.Background()
cfg := &config.Config{}
+20 -1
View File
@@ -136,7 +136,7 @@ func (s *TelegramBotService) generateCode(ctx context.Context, kind string, dura
// used only after the dependent action (account create / renew) succeeds, so a
// failed action never burns a code.
func (s *TelegramBotService) lookupRedeemableCode(ctx context.Context, raw, wantKind string) (*model.RegistrationCode, string) {
code := strings.TrimSpace(raw)
code := normalizeRedemptionCode(raw)
if code == "" {
return nil, "请提供兑换码。"
}
@@ -161,6 +161,25 @@ func (s *TelegramBotService) lookupRedeemableCode(ctx context.Context, raw, want
return rc, ""
}
func normalizeRedemptionCode(raw string) string {
code := strings.ToUpper(strings.TrimSpace(raw))
code = strings.NewReplacer(" ", "", "-", "", "_", "").Replace(code)
return code
}
func looksLikeRedemptionCode(raw string) bool {
code := normalizeRedemptionCode(raw)
if len(code) < 8 || len(code) > 32 {
return false
}
for _, ch := range code {
if !strings.ContainsRune(codeAlphabet, ch) {
return false
}
}
return true
}
// ── 续期 ────────────────────────────────────────────────────────────────────
// renewUser extends a user's expiry by durationDays. A nil/zero current expiry
+63
View File
@@ -2,6 +2,7 @@ package service
import (
"context"
"encoding/json"
"strings"
"testing"
"time"
@@ -552,6 +553,68 @@ func TestBotRedeemRegisterCodeCreatesOnlyOneAccount(t *testing.T) {
}
}
func TestBotRegisterCommandAcceptsRegistrationCode(t *testing.T) {
ctx := context.Background()
_, bot := newBotTestService(t)
code, err := bot.generateCode(ctx, model.RegistrationCodeRegister, 30, 0, "")
if err != nil {
t.Fatal(err)
}
channel := &model.NotifyChannel{Name: "Telegram", Type: "telegram", Enabled: true, Config: `{"admin_user_ids":"9301"}`}
msg := &TelegramMessage{From: TelegramUser{ID: 9301, Username: "codeuser"}, Chat: TelegramChat{ID: 9301, Type: "private"}}
reply, err := bot.executeCommand(ctx, channel, msg, "/register "+strings.ToLower(code.Code[:4])+"-"+strings.ToLower(code.Code[4:]))
if err != nil {
t.Fatal(err)
}
if !strings.Contains(reply.Text, "兑换成功") {
t.Fatalf("/register CODE should redeem registration code, got %q", reply.Text)
}
if binding := bot.telegramBinding(ctx, 9301); binding == nil {
t.Fatal("register code should bind the newly created account")
}
}
func TestBotPlainRegistrationCodeMessageRedeems(t *testing.T) {
ctx := context.Background()
repos, bot := newBotTestService(t)
code, err := bot.generateCode(ctx, model.RegistrationCodeRegister, 30, 0, "")
if err != nil {
t.Fatal(err)
}
if err := repos.DB.Create(&model.NotifyChannel{
Name: "Telegram",
Type: "telegram",
Enabled: true,
Config: `{"admin_user_ids":"9302"}`,
}).Error; err != nil {
t.Fatal(err)
}
update, _ := json.Marshal(TelegramUpdate{
UpdateID: 1,
Message: &TelegramMessage{
MessageID: 12,
Text: strings.ToLower(code.Code),
From: TelegramUser{ID: 9302, Username: "plaincode"},
Chat: TelegramChat{ID: 9302, Type: "private"},
},
})
if err := bot.HandleWebhook(ctx, update); err != nil {
t.Fatal(err)
}
if binding := bot.telegramBinding(ctx, 9302); binding == nil {
t.Fatal("plain code private message should redeem and bind account")
}
var used model.RegistrationCode
if err := repos.DB.Where("code = ?", code.Code).First(&used).Error; err != nil {
t.Fatal(err)
}
if used.UsedAt == nil || used.UsedByUserID == "" {
t.Fatal("plain code message should mark registration code as used")
}
}
func TestBotAdminCodeAndUserCommands(t *testing.T) {
ctx := context.Background()
repos, bot := newBotTestService(t)
+38 -9
View File
@@ -160,9 +160,12 @@ func (s *TelegramBotService) HandleWebhook(ctx context.Context, body []byte) err
if err := json.Unmarshal(body, &update); err != nil {
return fmt.Errorf("invalid update: %w", err)
}
return s.handleTelegramUpdate(ctx, update, nil)
}
func (s *TelegramBotService) handleTelegramUpdate(ctx context.Context, update TelegramUpdate, channelHint *model.NotifyChannel) error {
if update.CallbackQuery != nil {
return s.handleCallback(ctx, update.CallbackQuery)
return s.handleCallback(ctx, update.CallbackQuery, channelHint)
}
if update.Message == nil || update.Message.Text == "" {
@@ -177,7 +180,7 @@ func (s *TelegramBotService) HandleWebhook(ctx context.Context, body []byte) err
// command gate so the button-driven menu can collect free-form input.
if !telegramIsCommandText(text) {
if msg.Chat.Type == "" || msg.Chat.Type == "private" {
if channel := s.findChannelForMessage(ctx, msg); channel != nil {
if channel := s.channelForMessage(ctx, msg, channelHint); channel != nil {
if reply, handled := s.handlePendingText(ctx, channel, msg, text); handled {
if reply.Text != "" {
if err := s.reply(ctx, channel, msg.Chat.ID, reply); err != nil {
@@ -187,6 +190,16 @@ func (s *TelegramBotService) HandleWebhook(ctx context.Context, body []byte) err
s.deleteTelegramSourceMessage(channel, msg.Chat.ID, msg.MessageID)
return nil
}
if looksLikeRedemptionCode(text) {
reply := s.cmdRedeem(ctx, channel, msg, []string{text})
if reply.Text != "" {
if err := s.reply(ctx, channel, msg.Chat.ID, reply); err != nil {
s.log.Error("reply failed", zap.Error(err))
}
}
s.deleteTelegramSourceMessage(channel, msg.Chat.ID, msg.MessageID)
return nil
}
}
}
return nil
@@ -203,7 +216,7 @@ func (s *TelegramBotService) HandleWebhook(ctx context.Context, body []byte) err
// 获取该消息可使用的 Telegram 通知渠道配置。群组/频道消息必须来自
// 已配置的群组/频道;私聊消息会选择一个可验证该用户成员身份的 Bot。
channel := s.findChannelForMessage(ctx, msg)
channel := s.channelForMessage(ctx, msg, channelHint)
if channel == nil {
s.log.Warn("telegram channel not allowed or not configured",
zap.Int("chat_id", msg.Chat.ID),
@@ -360,6 +373,9 @@ func (s *TelegramBotService) cmdStart(ctx context.Context, msg *TelegramMessage,
// cmdRegister 处理 /register 命令:在管理员开启注册后,普通用户可通过 Bot
// 注册一个新的媒体中心账号,并自动绑定到当前 Telegram 账号。
func (s *TelegramBotService) cmdRegister(ctx context.Context, channel *model.NotifyChannel, msg *TelegramMessage, args []string) telegramCommandReply {
if len(args) == 1 && looksLikeRedemptionCode(args[0]) {
return s.redeemRegisterFlow(ctx, channel, msg, args[0])
}
if !s.openRegEnabled(ctx) {
return telegramCommandReply{Text: "注册功能未开放,请联系管理员开启后再试。"}
}
@@ -793,7 +809,8 @@ func (s *TelegramBotService) StartPolling(ctx context.Context) {
s.pollingCancel[botToken] = cancel
s.pollingMu.Unlock()
go s.pollLoop(pollCtx, cfg)
channel := ch
go s.pollLoop(pollCtx, cfg, &channel)
s.log.Info("started telegram polling", zap.String("channel", ch.Name))
}
}
@@ -810,7 +827,7 @@ func (s *TelegramBotService) StopPolling() {
}
// pollLoop 对单个 Bot Token 执行长轮询。
func (s *TelegramBotService) pollLoop(ctx context.Context, cfg map[string]string) {
func (s *TelegramBotService) pollLoop(ctx context.Context, cfg map[string]string, channel *model.NotifyChannel) {
var offset int64 = 0
pollURL, err := telegramMethodURL(cfg, cfg["bot_token"], "getUpdates")
if err != nil {
@@ -856,8 +873,7 @@ func (s *TelegramBotService) pollLoop(ctx context.Context, cfg map[string]string
go func(u TelegramUpdate) {
handlerCtx, cancel := context.WithTimeout(ctx, 2*time.Minute)
defer cancel()
raw, _ := json.Marshal(u)
_ = s.HandleWebhook(handlerCtx, raw)
_ = s.handleTelegramUpdate(handlerCtx, u, channel)
}(upd)
}
}
@@ -1086,13 +1102,26 @@ func (s *TelegramBotService) findChannelForMessage(ctx context.Context, msg *Tel
return first
}
func (s *TelegramBotService) handleCallback(ctx context.Context, cb *TelegramCallbackQuery) error {
func (s *TelegramBotService) channelForMessage(ctx context.Context, msg *TelegramMessage, hint *model.NotifyChannel) *model.NotifyChannel {
if hint == nil {
return s.findChannelForMessage(ctx, msg)
}
if msg == nil {
return hint
}
if msg.Chat.Type != "" && msg.Chat.Type != "private" && !s.telegramChatAllowed(hint, msg.Chat.ID) {
return nil
}
return hint
}
func (s *TelegramBotService) handleCallback(ctx context.Context, cb *TelegramCallbackQuery, channelHint *model.NotifyChannel) error {
if cb == nil || cb.Message == nil {
return nil
}
msg := *cb.Message
msg.From = cb.From
channel := s.findChannelForMessage(ctx, &msg)
channel := s.channelForMessage(ctx, &msg, channelHint)
if channel == nil {
channel = s.findChannelByChatID(ctx, cb.Message.Chat.ID)
}
@@ -191,6 +191,62 @@ func TestTelegramGroupAdminMenuDoesNotExposeButtonsInGroup(t *testing.T) {
}
}
func TestTelegramPollingChannelHintWinsForPrivateMessages(t *testing.T) {
ctx := t.Context()
repos, bot := newBotTestService(t)
msg := &TelegramMessage{
From: TelegramUser{ID: 9101, Username: "viewer", FirstName: "Viewer"},
Chat: TelegramChat{ID: 9101, Type: "private"},
}
bad := model.NotifyChannel{Name: "BadToken", Type: "telegram", Enabled: true, Config: `{"bot_token":"bad","admin_user_ids":"9101"}`}
good := model.NotifyChannel{Name: "GoodToken", Type: "telegram", Enabled: true, Config: `{"bot_token":"good","admin_user_ids":"9101"}`}
if err := repos.DB.Create(&bad).Error; err != nil {
t.Fatal(err)
}
if err := repos.DB.Create(&good).Error; err != nil {
t.Fatal(err)
}
if first := bot.findChannelForMessage(ctx, msg); first == nil || first.ID != bad.ID {
t.Fatalf("setup expected normal private lookup to pick first channel, got %#v", first)
}
if hinted := bot.channelForMessage(ctx, msg, &good); hinted == nil || hinted.ID != good.ID {
t.Fatalf("polling channel hint should route replies through the token that received the update, got %#v", hinted)
}
}
func TestTelegramSakuraCompatibleUserCommands(t *testing.T) {
ctx := t.Context()
repos, bot := newBotTestService(t)
user := &model.User{Username: "viewer", PasswordHash: "hash", Role: "user", IsActive: true}
if err := repos.User.Create(ctx, user); err != nil {
t.Fatal(err)
}
if err := repos.DB.Create(&model.TelegramBinding{TelegramUserID: 9102, ChatID: 9102, UserID: user.ID}).Error; err != nil {
t.Fatal(err)
}
channel := &model.NotifyChannel{Name: "Telegram", Type: "telegram", Enabled: true, Config: `{"admin_user_ids":"9102"}`}
msg := &TelegramMessage{
From: TelegramUser{ID: 9102, Username: "viewer", FirstName: "Viewer"},
Chat: TelegramChat{ID: 9102, Type: "private"},
}
info, err := bot.executeCommand(ctx, channel, msg, "/myinfo")
if err != nil {
t.Fatal(err)
}
if !strings.Contains(info.Text, "我的账号") {
t.Fatalf("/myinfo should show account info, got %q", info.Text)
}
count, err := bot.executeCommand(ctx, channel, msg, "/count")
if err != nil {
t.Fatal(err)
}
if !strings.Contains(count.Text, "媒体库统计") {
t.Fatalf("/count should show library counts, got %q", count.Text)
}
}
func telegramReplyHasButtonPrefix(reply telegramCommandReply, prefix string) bool {
for _, row := range reply.Buttons {
for _, button := range row {
+5 -2
View File
@@ -40,7 +40,8 @@ func (s *TelegramBotService) telegramCommandDefinitions(ctx context.Context, cha
return telegramCommandReply{Text: s.cmdHelp(ctx, msg)}, nil
}},
{Aliases: []string{"/hideadult", "/hide_adult", "/adult"}, GroupAllowed: true, Handle: func(args []string) (telegramCommandReply, error) { return s.cmdHideAdult(ctx, msg, args), nil }},
{Aliases: []string{"/account", "/me"}, GroupAllowed: true, Handle: func(args []string) (telegramCommandReply, error) { return s.replyAccount(ctx, msg), nil }},
{Aliases: []string{"/account", "/me", "/myinfo"}, GroupAllowed: true, Handle: func(args []string) (telegramCommandReply, error) { return s.replyAccount(ctx, msg), nil }},
{Aliases: []string{"/count"}, GroupAllowed: true, Handle: func(args []string) (telegramCommandReply, error) { return s.cmdStats(ctx) }},
{Aliases: []string{"/signin", "/checkin"}, GroupAllowed: true, Handle: func(args []string) (telegramCommandReply, error) { return s.replySignIn(ctx, msg), nil }},
{Aliases: []string{"/devices"}, GroupAllowed: true, Handle: func(args []string) (telegramCommandReply, error) { return s.replyDevices(ctx, msg), nil }},
{Aliases: []string{"/kick"}, GroupAllowed: true, Handle: func(args []string) (telegramCommandReply, error) { return s.cmdKick(ctx, msg, args), nil }},
@@ -128,7 +129,7 @@ func telegramSupportedCommand(cmd string) bool {
var telegramSupportedCommandSet = map[string]struct{}{
"/start": {}, "/menu": {}, "/cancel": {}, "/help": {}, "/hideadult": {}, "/hide_adult": {}, "/adult": {},
"/account": {}, "/me": {}, "/signin": {}, "/checkin": {}, "/devices": {}, "/kick": {}, "/setname": {}, "/rename": {}, "/setpass": {}, "/passwd": {}, "/password": {},
"/account": {}, "/me": {}, "/myinfo": {}, "/count": {}, "/signin": {}, "/checkin": {}, "/devices": {}, "/kick": {}, "/setname": {}, "/rename": {}, "/setpass": {}, "/passwd": {}, "/password": {},
"/redeem": {}, "/redeem_register": {}, "/redeem_renew": {},
"/register": {}, "/reg": {}, "/signup": {}, "/registration": {}, "/reg_switch": {}, "/openreg": {},
"/capacity": {}, "/users": {}, "/gencode": {}, "/renew_user": {}, "/delete_user": {}, "/unbind": {}, "/unbind_duplicates": {}, "/unbind_inactive": {},
@@ -151,6 +152,8 @@ func telegramPrivateBotCommandMenu() []telegramBotCommand {
{Command: "menu", Description: "打开功能菜单"},
{Command: "help", Description: "查看命令帮助"},
{Command: "account", Description: "查看账号状态"},
{Command: "myinfo", Description: "查看账号状态"},
{Command: "count", Description: "查看媒体库数量"},
{Command: "signin", Description: "签到"},
{Command: "devices", Description: "查看登录设备"},
{Command: "kick", Description: "踢下线设备"},
+1 -1
View File
@@ -549,7 +549,7 @@ func (s *TelegramBotService) redeemRegisterFlow(ctx context.Context, channel *mo
}
func (s *TelegramBotService) createUserFromRegistrationCode(ctx context.Context, rawCode string) (*model.User, string, *model.RegistrationCode, error) {
code := strings.TrimSpace(rawCode)
code := normalizeRedemptionCode(rawCode)
if code == "" {
return nil, "", nil, errRegistrationCodeAlreadyUsed
}