mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-25 18:22:41 +08:00
Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c86867ff8b | ||
|
|
7c4c7ef3c7 | ||
|
|
21a519fdbe | ||
|
|
1d431dbc88 |
@@ -1,10 +1,13 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
"github.com/celestix/gotgproto/dispatcher"
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/duke-git/lancet/v2/slice"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/dirutil"
|
||||
"github.com/krau/SaveAny-Bot/client/bot/handlers/utils/msgelem"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
@@ -12,10 +15,23 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
)
|
||||
|
||||
// responsibleUserID returns the sender's ID. Callback queries carry it
|
||||
// natively; message updates resolve it through the entity map.
|
||||
func responsibleUserID(u *ext.Update) int64 {
|
||||
if u.CallbackQuery != nil {
|
||||
return u.CallbackQuery.GetUserID()
|
||||
}
|
||||
return u.GetUserChat().GetID()
|
||||
}
|
||||
|
||||
func checkPermission(ctx *ext.Context, update *ext.Update) error {
|
||||
userID := update.GetUserChat().GetID()
|
||||
userID := responsibleUserID(update)
|
||||
if !slice.Contain(config.C().GetUsersID(), userID) {
|
||||
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgCommonErrorNoPermission, nil)), nil)
|
||||
if cbq := update.CallbackQuery; cbq != nil {
|
||||
ctx.AnswerCallback(msgelem.AlertCallbackAnswer(cbq.GetQueryID(), i18n.T(i18nk.BotMsgCommonErrorNoPermission, nil)))
|
||||
} else {
|
||||
ctx.Reply(update, ext.ReplyTextString(i18n.T(i18nk.BotMsgCommonErrorNoPermission, nil)), nil)
|
||||
}
|
||||
return dispatcher.EndGroups
|
||||
}
|
||||
|
||||
@@ -23,10 +39,12 @@ func checkPermission(ctx *ext.Context, update *ext.Update) error {
|
||||
}
|
||||
|
||||
// withPermission wraps a callback handler with the same whitelist check used
|
||||
// for message handlers (checkPermission).
|
||||
// for message handlers (checkPermission). ContinueGroups is the dispatcher's
|
||||
// success sentinel, not an error: only real failures and EndGroups stop the
|
||||
// chain before the wrapped handler runs.
|
||||
func withPermission(handler func(*ext.Context, *ext.Update) error) func(*ext.Context, *ext.Update) error {
|
||||
return func(ctx *ext.Context, update *ext.Update) error {
|
||||
if err := checkPermission(ctx, update); err != nil {
|
||||
if err := checkPermission(ctx, update); err != nil && !errors.Is(err, dispatcher.ContinueGroups) {
|
||||
return err
|
||||
}
|
||||
return handler(ctx, update)
|
||||
|
||||
85
client/bot/handlers/middleware_test.go
Normal file
85
client/bot/handlers/middleware_test.go
Normal file
@@ -0,0 +1,85 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/celestix/gotgproto/ext"
|
||||
"github.com/celestix/gotgproto/types"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
)
|
||||
|
||||
// Regression: callback queries usually arrive as updateShort without entity
|
||||
// maps, so resolving the sender through the entity map yields ID 0 and every
|
||||
// click was denied by the whitelist check. Callback updates must use the
|
||||
// native UserID field.
|
||||
func TestResponsibleUserID(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
update *ext.Update
|
||||
want int64
|
||||
}{
|
||||
{
|
||||
name: "callback query uses native user id",
|
||||
update: &ext.Update{CallbackQuery: &tg.UpdateBotCallbackQuery{UserID: 42}},
|
||||
want: 42,
|
||||
},
|
||||
{
|
||||
name: "message resolves through entity map",
|
||||
update: &ext.Update{
|
||||
EffectiveMessage: &types.Message{Message: &tg.Message{PeerID: &tg.PeerUser{UserID: 7}}},
|
||||
Entities: &tg.Entities{Users: map[int64]*tg.User{7: {ID: 7}}},
|
||||
},
|
||||
want: 7,
|
||||
},
|
||||
{
|
||||
name: "callback query ignores entity map",
|
||||
update: &ext.Update{
|
||||
CallbackQuery: &tg.UpdateBotCallbackQuery{UserID: 9},
|
||||
Entities: &tg.Entities{Users: map[int64]*tg.User{8: {ID: 8}}},
|
||||
},
|
||||
want: 9,
|
||||
},
|
||||
{
|
||||
name: "unresolvable update yields zero",
|
||||
update: &ext.Update{},
|
||||
want: 0,
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := responsibleUserID(tt.update); got != tt.want {
|
||||
t.Fatalf("responsibleUserID() = %d, want %d", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Regression: withPermission must treat ContinueGroups (the dispatcher's
|
||||
// success sentinel) as a pass and invoke the wrapped handler. v0.60.1 treated
|
||||
// it as an error, so every permitted callback was swallowed before the real
|
||||
// handler ran.
|
||||
func TestWithPermissionInvokesHandler(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "config.toml")
|
||||
if err := os.WriteFile(path, []byte("workers = 2\n\n[[users]]\nid = 42\n"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := config.Init(t.Context(), path); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
update := &ext.Update{CallbackQuery: &tg.UpdateBotCallbackQuery{UserID: 42}}
|
||||
called := false
|
||||
handler := withPermission(func(ctx *ext.Context, u *ext.Update) error {
|
||||
called = true
|
||||
return nil
|
||||
})
|
||||
if err := handler(&ext.Context{}, update); err != nil {
|
||||
t.Fatalf("withPermission returned error: %v", err)
|
||||
}
|
||||
if !called {
|
||||
t.Fatal("withPermission did not invoke the wrapped handler")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user