From c2cfa5c0150f55e623dd525cda916f7d503522e6 Mon Sep 17 00:00:00 2001 From: biast12 Date: Tue, 10 Mar 2026 18:33:46 +0100 Subject: [PATCH] Handle worker context --- .../admindebug/server/blacklistreason.go | 8 ++++- .../admindebug/server/entitlements.go | 8 ++++- .../admindebug/server/modals/permissions.go | 31 ++-------------- .../server/modals/ticketpermissions.go | 23 +++++++----- .../admindebug/server/modals/usertickets.go | 13 +++++-- .../admindebug/server/monitoredbots.go | 8 ++++- .../handlers/admindebug/server/recache.go | 31 +--------------- bot/command/impl/admin/adminblacklist.go | 26 ++------------ bot/command/impl/admin/adminrecache.go | 30 +--------------- bot/utils/context.go | 35 +++++++++++++++++++ 10 files changed, 89 insertions(+), 124 deletions(-) diff --git a/bot/button/handlers/admindebug/server/blacklistreason.go b/bot/button/handlers/admindebug/server/blacklistreason.go index 8e80a55..f399a59 100644 --- a/bot/button/handlers/admindebug/server/blacklistreason.go +++ b/bot/button/handlers/admindebug/server/blacklistreason.go @@ -43,8 +43,14 @@ func (h *AdminDebugServerBlacklistReasonHandler) Execute(ctx *context.ButtonCont return } + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) + if err != nil { + ctx.HandleError(err) + return + } + // Get guild to fetch owner ID - guild, err := ctx.Worker().GetGuild(guildId) + guild, err := worker.GetGuild(guildId) if err != nil { ctx.HandleError(err) return diff --git a/bot/button/handlers/admindebug/server/entitlements.go b/bot/button/handlers/admindebug/server/entitlements.go index f20f178..f4c8cab 100644 --- a/bot/button/handlers/admindebug/server/entitlements.go +++ b/bot/button/handlers/admindebug/server/entitlements.go @@ -42,8 +42,14 @@ func (h *AdminDebugServerEntitlementsHandler) Execute(ctx *context.ButtonContext return } + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) + if err != nil { + ctx.HandleError(err) + return + } + // Get guild to fetch owner ID - guild, err := ctx.Worker().GetGuild(guildId) + guild, err := worker.GetGuild(guildId) if err != nil { ctx.HandleError(err) return diff --git a/bot/button/handlers/admindebug/server/modals/permissions.go b/bot/button/handlers/admindebug/server/modals/permissions.go index 1bd2609..3417e56 100644 --- a/bot/button/handlers/admindebug/server/modals/permissions.go +++ b/bot/button/handlers/admindebug/server/modals/permissions.go @@ -82,38 +82,13 @@ func (h *AdminDebugServerPermissionsModalSubmitHandler) Execute(ctx *context.Mod selectedValues := selectData.Values - // Get guild and settings - botId, isWhitelabel, err := dbclient.Client.WhitelabelGuilds.GetBotByGuild(ctx, guildId) + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) if err != nil { ctx.HandleError(err) return } - var worker *w.Context - if isWhitelabel { - bot, err := dbclient.Client.Whitelabel.GetByBotId(ctx, botId) - if err != nil { - ctx.HandleError(err) - return - } - - if bot.BotId == 0 { - ctx.HandleError(errors.New("bot not found")) - return - } - - worker = &w.Context{ - Token: bot.Token, - BotId: bot.BotId, - IsWhitelabel: true, - ShardId: 0, - Cache: ctx.Worker().Cache, - RateLimiter: nil, - } - } else { - worker = ctx.Worker() - } - + // Get guild and settings settings, err := dbclient.Client.Settings.Get(ctx, guildId) if err != nil { ctx.HandleError(err) @@ -126,7 +101,7 @@ func (h *AdminDebugServerPermissionsModalSubmitHandler) Execute(ctx *context.Mod return } - botMember, err := worker.GetGuildMember(guildId, ctx.Worker().BotId) + botMember, err := worker.GetGuildMember(guildId, worker.BotId) if err != nil { ctx.HandleError(err) return diff --git a/bot/button/handlers/admindebug/server/modals/ticketpermissions.go b/bot/button/handlers/admindebug/server/modals/ticketpermissions.go index 28de26c..fe9b5f7 100644 --- a/bot/button/handlers/admindebug/server/modals/ticketpermissions.go +++ b/bot/button/handlers/admindebug/server/modals/ticketpermissions.go @@ -11,6 +11,7 @@ import ( "github.com/TicketsBot-cloud/database" "github.com/TicketsBot-cloud/gdl/objects/interaction" "github.com/TicketsBot-cloud/gdl/objects/interaction/component" + w "github.com/TicketsBot-cloud/worker" "github.com/TicketsBot-cloud/worker/bot/button/registry" "github.com/TicketsBot-cloud/worker/bot/button/registry/matcher" "github.com/TicketsBot-cloud/worker/bot/command" @@ -94,8 +95,14 @@ func (h *AdminDebugServerTicketPermissionsModalSubmitHandler) Execute(ctx *conte return } + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) + if err != nil { + ctx.HandleError(err) + return + } + // Get guild info - guild, err := ctx.Worker().GetGuild(guildId) + guild, err := worker.GetGuild(guildId) if err != nil { ctx.HandleError(err) return @@ -150,10 +157,10 @@ func (h *AdminDebugServerTicketPermissionsModalSubmitHandler) Execute(ctx *conte } if isUser { - result := checkUserTicketPermissions(ctx, guildId, guild.OwnerId, entityId, adminUsers, supportUsers, panels) + result := checkUserTicketPermissions(ctx, worker, guildId, guild.OwnerId, entityId, adminUsers, supportUsers, panels) results = append(results, result) } else { - result := checkRoleTicketPermissions(ctx, guildId, entityId, adminRoles, supportRoles, panels) + result := checkRoleTicketPermissions(ctx, worker, guildId, entityId, adminRoles, supportRoles, panels) results = append(results, result) } } @@ -173,7 +180,7 @@ func (h *AdminDebugServerTicketPermissionsModalSubmitHandler) Execute(ctx *conte })) } -func checkUserTicketPermissions(ctx *context.ModalContext, guildId, ownerId, userId uint64, adminUsers, supportUsers []uint64, panels []database.Panel) string { +func checkUserTicketPermissions(ctx *context.ModalContext, worker *w.Context, guildId, ownerId, userId uint64, adminUsers, supportUsers []uint64, panels []database.Panel) string { var lines []string // Ticket Permission Level @@ -190,10 +197,10 @@ func checkUserTicketPermissions(ctx *context.ModalContext, guildId, ownerId, use lines = append(lines, ticketPermLevel) // Check member roles for role-based permissions - member, err := ctx.Worker().GetGuildMember(guildId, userId) + member, err := worker.GetGuildMember(guildId, userId) if err == nil { // Get all guild roles for name lookups - guildRoles, err := ctx.Worker().GetGuildRoles(guildId) + guildRoles, err := worker.GetGuildRoles(guildId) roleMap := make(map[uint64]string) if err == nil { for _, role := range guildRoles { @@ -300,11 +307,11 @@ func checkUserTicketPermissions(ctx *context.ModalContext, guildId, ownerId, use return fmt.Sprintf("**User:** <@%d>\n%s", userId, strings.Join(lines, "\n")) } -func checkRoleTicketPermissions(ctx *context.ModalContext, guildId, roleId uint64, adminRoles, supportRoles []uint64, panels []database.Panel) string { +func checkRoleTicketPermissions(ctx *context.ModalContext, worker *w.Context, guildId, roleId uint64, adminRoles, supportRoles []uint64, panels []database.Panel) string { var lines []string // Get role name - guildRoles, err := ctx.Worker().GetGuildRoles(guildId) + guildRoles, err := worker.GetGuildRoles(guildId) roleName := "Unknown Role" if err == nil { for _, role := range guildRoles { diff --git a/bot/button/handlers/admindebug/server/modals/usertickets.go b/bot/button/handlers/admindebug/server/modals/usertickets.go index 28433b8..98b9530 100644 --- a/bot/button/handlers/admindebug/server/modals/usertickets.go +++ b/bot/button/handlers/admindebug/server/modals/usertickets.go @@ -11,6 +11,7 @@ import ( "github.com/TicketsBot-cloud/database" "github.com/TicketsBot-cloud/gdl/objects/interaction" "github.com/TicketsBot-cloud/gdl/objects/interaction/component" + w "github.com/TicketsBot-cloud/worker" "github.com/TicketsBot-cloud/worker/bot/button/registry" "github.com/TicketsBot-cloud/worker/bot/button/registry/matcher" "github.com/TicketsBot-cloud/worker/bot/command" @@ -93,6 +94,12 @@ func (h *AdminDebugServerUserTicketsModalSubmitHandler) Execute(ctx *context.Mod return } + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) + if err != nil { + ctx.HandleError(err) + return + } + // Get all open tickets for the guild allOpenTickets, err := dbclient.Client.Tickets.GetGuildOpenTickets(ctx, guildId) if err != nil { @@ -104,7 +111,7 @@ func (h *AdminDebugServerUserTicketsModalSubmitHandler) Execute(ctx *context.Mod var results []string for _, userId := range userIds { - result := checkUserTickets(ctx, guildId, userId, allOpenTickets) + result := checkUserTickets(ctx, worker, guildId, userId, allOpenTickets) results = append(results, result) } @@ -123,7 +130,7 @@ func (h *AdminDebugServerUserTicketsModalSubmitHandler) Execute(ctx *context.Mod })) } -func checkUserTickets(ctx *context.ModalContext, guildId, userId uint64, allOpenTickets []database.Ticket) string { +func checkUserTickets(ctx *context.ModalContext, worker *w.Context, guildId, userId uint64, allOpenTickets []database.Ticket) string { var lines []string // Get open ticket count @@ -162,7 +169,7 @@ func checkUserTickets(ctx *context.ModalContext, guildId, userId uint64, allOpen channelMention := fmt.Sprintf("`%d`", *ticket.ChannelId) if ticket.ChannelId != nil { - _, err := ctx.Worker().GetChannel(*ticket.ChannelId) + _, err := worker.GetChannel(*ticket.ChannelId) if err == nil { channelExists = true } diff --git a/bot/button/handlers/admindebug/server/monitoredbots.go b/bot/button/handlers/admindebug/server/monitoredbots.go index 86c204b..614acc2 100644 --- a/bot/button/handlers/admindebug/server/monitoredbots.go +++ b/bot/button/handlers/admindebug/server/monitoredbots.go @@ -41,10 +41,16 @@ func (h *AdminDebugServerMonitoredBotsHandler) Execute(ctx *context.ButtonContex return } + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) + if err != nil { + ctx.HandleError(err) + return + } + var monitoredBotsPresent []string for _, botId := range config.Conf.Bot.MonitoredBots { - _, err = ctx.Worker().GetGuildMember(guildId, botId) + _, err = worker.GetGuildMember(guildId, botId) if err == nil { botUser, err := ctx.Worker().GetUser(botId) if err == nil { diff --git a/bot/button/handlers/admindebug/server/recache.go b/bot/button/handlers/admindebug/server/recache.go index 35da4ba..1f1f2a7 100644 --- a/bot/button/handlers/admindebug/server/recache.go +++ b/bot/button/handlers/admindebug/server/recache.go @@ -1,7 +1,6 @@ package server import ( - "errors" "fmt" "strconv" "strings" @@ -9,13 +8,11 @@ import ( permcache "github.com/TicketsBot-cloud/common/permission" "github.com/TicketsBot-cloud/gdl/objects/interaction/component" - w "github.com/TicketsBot-cloud/worker" "github.com/TicketsBot-cloud/worker/bot/button/registry" "github.com/TicketsBot-cloud/worker/bot/button/registry/matcher" "github.com/TicketsBot-cloud/worker/bot/command" "github.com/TicketsBot-cloud/worker/bot/command/context" "github.com/TicketsBot-cloud/worker/bot/customisation" - "github.com/TicketsBot-cloud/worker/bot/dbclient" "github.com/TicketsBot-cloud/worker/bot/redis" "github.com/TicketsBot-cloud/worker/bot/utils" ) @@ -38,7 +35,6 @@ func (h *AdminDebugServerRecacheHandler) Properties() registry.Properties { } func (h *AdminDebugServerRecacheHandler) Execute(ctx *context.ButtonContext) { - if !utils.IsBotHelper(ctx.UserId()) { ctx.ReplyRaw(customisation.Red, "Error", "You do not have permission to use this button.") } @@ -66,37 +62,12 @@ func (h *AdminDebugServerRecacheHandler) Execute(ctx *context.ButtonContext) { ctx.Worker().Cache.DeleteGuildChannels(ctx, guildId) ctx.Worker().Cache.DeleteGuildRoles(ctx, guildId) - botId, isWhitelabel, err := dbclient.Client.WhitelabelGuilds.GetBotByGuild(ctx, guildId) + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) if err != nil { ctx.HandleError(err) return } - var worker *w.Context - if isWhitelabel { - bot, err := dbclient.Client.Whitelabel.GetByBotId(ctx, botId) - if err != nil { - ctx.HandleError(err) - return - } - - if bot.BotId == 0 { - ctx.HandleError(errors.New("bot not found")) - return - } - - worker = &w.Context{ - Token: bot.Token, - BotId: bot.BotId, - IsWhitelabel: true, - ShardId: 0, - Cache: ctx.Worker().Cache, - RateLimiter: nil, // Use http-proxy ratelimit functionality - } - } else { - worker = ctx.Worker() - } - guild, err := worker.GetGuild(guildId) if err != nil { ctx.HandleError(err) diff --git a/bot/command/impl/admin/adminblacklist.go b/bot/command/impl/admin/adminblacklist.go index 9ebe521..ec16857 100644 --- a/bot/command/impl/admin/adminblacklist.go +++ b/bot/command/impl/admin/adminblacklist.go @@ -9,7 +9,6 @@ import ( "github.com/TicketsBot-cloud/common/permission" "github.com/TicketsBot-cloud/gdl/objects/interaction" "github.com/TicketsBot-cloud/gdl/objects/interaction/component" - "github.com/TicketsBot-cloud/worker" "github.com/TicketsBot-cloud/worker/bot/blacklist" "github.com/TicketsBot-cloud/worker/bot/command" "github.com/TicketsBot-cloud/worker/bot/command/registry" @@ -75,35 +74,16 @@ func (AdminBlacklistCommand) Execute(ctx registry.CommandContext, guildIdRaw str } // Check for whitelabel - botId, isWhitelabel, err := dbclient.Client.WhitelabelGuilds.GetBotByGuild(ctx, guildId) + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) if err != nil { ctx.HandleError(err) return } - var w *worker.Context - if isWhitelabel { - bot, err := dbclient.Client.Whitelabel.GetByBotId(ctx, botId) - if err != nil { - ctx.HandleError(err) - return - } - - w = &worker.Context{ - Token: bot.Token, - BotId: bot.BotId, - IsWhitelabel: true, - Cache: ctx.Worker().Cache, - RateLimiter: nil, // Use http-proxy ratelimit functionality - } - } else { - w = ctx.Worker() - } - // Try to get owner var ownerId *uint64 var botInGuild bool - guild, err := w.GetGuild(guildId) + guild, err := worker.GetGuild(guildId) if err == nil { ownerId = &guild.OwnerId botInGuild = true @@ -142,7 +122,7 @@ func (AdminBlacklistCommand) Execute(ctx registry.CommandContext, guildIdRaw str // Leave guild if botInGuild { - if err := w.LeaveGuild(guildId); err != nil { + if err := worker.LeaveGuild(guildId); err != nil { ctx.HandleError(err) return } diff --git a/bot/command/impl/admin/adminrecache.go b/bot/command/impl/admin/adminrecache.go index 05809ee..fd0a124 100644 --- a/bot/command/impl/admin/adminrecache.go +++ b/bot/command/impl/admin/adminrecache.go @@ -1,7 +1,6 @@ package admin import ( - "errors" "fmt" "strconv" "time" @@ -9,11 +8,9 @@ import ( "github.com/TicketsBot-cloud/common/permission" "github.com/TicketsBot-cloud/gdl/objects/interaction" "github.com/TicketsBot-cloud/gdl/objects/interaction/component" - w "github.com/TicketsBot-cloud/worker" "github.com/TicketsBot-cloud/worker/bot/command" "github.com/TicketsBot-cloud/worker/bot/command/registry" "github.com/TicketsBot-cloud/worker/bot/customisation" - "github.com/TicketsBot-cloud/worker/bot/dbclient" "github.com/TicketsBot-cloud/worker/bot/redis" "github.com/TicketsBot-cloud/worker/bot/utils" "github.com/TicketsBot-cloud/worker/i18n" @@ -77,37 +74,12 @@ func (AdminRecacheCommand) Execute(ctx registry.CommandContext, providedGuildId ctx.Worker().Cache.DeleteGuildRoles(ctx, guildId) // re-cache - botId, isWhitelabel, err := dbclient.Client.WhitelabelGuilds.GetBotByGuild(ctx, guildId) + worker, err := utils.WorkerForGuild(ctx, ctx.Worker(), guildId) if err != nil { ctx.HandleError(err) return } - var worker *w.Context - if isWhitelabel { - bot, err := dbclient.Client.Whitelabel.GetByBotId(ctx, botId) - if err != nil { - ctx.HandleError(err) - return - } - - if bot.BotId == 0 { - ctx.HandleError(errors.New("bot not found")) - return - } - - worker = &w.Context{ - Token: bot.Token, - BotId: bot.BotId, - IsWhitelabel: true, - ShardId: 0, - Cache: ctx.Worker().Cache, - RateLimiter: nil, // Use http-proxy ratelimit functionality - } - } else { - worker = ctx.Worker() - } - guild, err := worker.GetGuild(guildId) if err != nil { ctx.HandleError(err) diff --git a/bot/utils/context.go b/bot/utils/context.go index f5acb8d..e278728 100644 --- a/bot/utils/context.go +++ b/bot/utils/context.go @@ -2,9 +2,44 @@ package utils import ( "context" + "errors" "time" + + w "github.com/TicketsBot-cloud/worker" + "github.com/TicketsBot-cloud/worker/bot/dbclient" ) func ContextTimeout(timeout time.Duration) (context.Context, context.CancelFunc) { return context.WithTimeout(context.Background(), timeout) } + +// WorkerForGuild returns a worker for the given guild. +// For whitelabel guilds it uses the whitelabel bot's token; otherwise it returns mainWorker. +func WorkerForGuild(ctx context.Context, mainWorker *w.Context, guildId uint64) (*w.Context, error) { + botId, isWhitelabel, err := dbclient.Client.WhitelabelGuilds.GetBotByGuild(ctx, guildId) + if err != nil { + return nil, err + } + + if !isWhitelabel { + return mainWorker, nil + } + + bot, err := dbclient.Client.Whitelabel.GetByBotId(ctx, botId) + if err != nil { + return nil, err + } + + if bot.BotId == 0 { + return nil, errors.New("bot not found") + } + + return &w.Context{ + Token: bot.Token, + BotId: bot.BotId, + IsWhitelabel: true, + ShardId: 0, + Cache: mainWorker.Cache, + RateLimiter: nil, // Use http-proxy ratelimit functionality + }, nil +}