From d7e44b556ec7fd1b4b5afcc8f78ee1878d26a130 Mon Sep 17 00:00:00 2001 From: AH-dark Date: Fri, 1 Mar 2024 23:30:35 +0800 Subject: [PATCH] feat: datacenter query --- components/basic-handler/internal/bot/bind.go | 5 + .../basic-handler/internal/bot/handlers/dc.go | 111 ++++++++++++++++++ .../internal/bot/handlers/dc.tpl | 13 ++ .../internal/bot/handlers/handlers.go | 6 + components/basic-handler/services/0module.go | 5 +- .../services/datacenter/service.go | 79 +++++++++++++ .../services/datacenter/service_test.go | 22 ++++ 7 files changed, 240 insertions(+), 1 deletion(-) create mode 100644 components/basic-handler/internal/bot/handlers/dc.go create mode 100644 components/basic-handler/internal/bot/handlers/dc.tpl create mode 100644 components/basic-handler/services/datacenter/service.go create mode 100644 components/basic-handler/services/datacenter/service_test.go diff --git a/components/basic-handler/internal/bot/bind.go b/components/basic-handler/internal/bot/bind.go index bc3d41c..e92b2b9 100644 --- a/components/basic-handler/internal/bot/bind.go +++ b/components/basic-handler/internal/bot/bind.go @@ -13,6 +13,11 @@ func BindHandlers(r *telegohandler.BotHandler, handlers handlers.Handlers) { // id command r.Handle(handlers.IDCommandHandler, telegohandler.CommandEqual("id")) + // datacenter command + r.Handle(handlers.DatacenterCommandHandler, telegohandler.CommandEqual("datacenter")) + r.Handle(handlers.DatacenterCommandHandler, telegohandler.CommandEqual("dc")) + r.Handle(handlers.DatacenterMoreInfoHandler, telegohandler.CallbackDataEqual("datacenter_more_info")) + // action command r.Handle(handlers.ActionCommandHandler, telegohandler.AnyMessageWithText(), telegohandler.TextPrefix("/"), telegohandler.Not(utils.PrivateChatOnly())) } diff --git a/components/basic-handler/internal/bot/handlers/dc.go b/components/basic-handler/internal/bot/handlers/dc.go new file mode 100644 index 0000000..ca5477b --- /dev/null +++ b/components/basic-handler/internal/bot/handlers/dc.go @@ -0,0 +1,111 @@ +package handlers + +import ( + _ "embed" + "fmt" + "html/template" + + "github.com/cloudwego/hertz/pkg/common/bytebufferpool" + "github.com/go-redis/redis_rate/v10" + "github.com/mymmrac/telego" + "github.com/mymmrac/telego/telegoutil" + "github.com/uptrace/opentelemetry-go-extra/otelzap" + "go.uber.org/zap" +) + +//go:embed dc.tpl +var dcTemplateText string + +func (h *handlers) DatacenterCommandHandler(bot *telego.Bot, update telego.Update) { + ctx, span := tracer.Start(update.Context(), "handlers.DatacenterCommandHandler") + defer span.End() + + if update.Message == nil { + otelzap.L().Ctx(ctx).Warn("update message is nil") + return + } + + if res, err := h.RedisRateLimiter.Allow(ctx, fmt.Sprintf("rate:basic-handler:handler:datacenter_command_handler:chat-%d", update.Message.Chat.ID), redis_rate.PerSecond(1)); err != nil { + otelzap.L().Ctx(ctx).Error("rate limit exceeded", zap.Error(err)) + _, _ = bot.SendMessage(telegoutil.Message(telegoutil.ID(update.Message.Chat.ID), "rate limit exceeded").WithReplyParameters(&telego.ReplyParameters{MessageID: update.Message.MessageID})) + return + } else if res.Allowed == 0 { + otelzap.L().Ctx(ctx).Warn("rate limit exceeded") + return + } + + var funcMap = template.FuncMap{ + "datacenter": func(username string) int { + dc, err := h.DatacenterService.QueryDatacenterByUsername(ctx, username) + if err != nil { + otelzap.L().Ctx(ctx).Error("failed to query datacenter", zap.Error(err)) + return 0 + } + + return dc + }, + } + + tpl, err := template.New("datacenter").Funcs(funcMap).Parse(dcTemplateText) + if err != nil { + otelzap.L().Ctx(ctx).Error("failed to parse template", zap.Error(err)) + return + } + + buf := bytebufferpool.Get() + defer bytebufferpool.Put(buf) + + if err := tpl.Execute(buf, update); err != nil { + otelzap.L().Ctx(ctx).Error("failed to execute template", zap.Error(err)) + return + } + + if _, err := bot.SendMessage( + telegoutil. + Message(telegoutil.ID(update.Message.Chat.ID), buf.String()). + WithParseMode(telego.ModeHTML). + WithReplyParameters(&telego.ReplyParameters{MessageID: update.Message.MessageID}). + WithReplyMarkup(telegoutil.InlineKeyboard( + telegoutil.InlineKeyboardRow( + telegoutil.InlineKeyboardButton("更多信息").WithCallbackData("datacenter_more_info"), + ), + )), + ); err != nil { + otelzap.L().Ctx(ctx).Error("failed to send message", zap.Error(err)) + return + } +} + +func (h *handlers) DatacenterMoreInfoHandler(bot *telego.Bot, update telego.Update) { + ctx, span := tracer.Start(update.Context(), "handlers.DatacenterMoreInfoHandler") + defer span.End() + + if !update.CallbackQuery.Message.IsAccessible() { + otelzap.L().Ctx(ctx).Warn("message is not accessible") + return + } + + chat := update.CallbackQuery.Message.GetChat() + if res, err := h.RedisRateLimiter.Allow(ctx, fmt.Sprintf("rate:basic-handler:handler:datacenter_more_info:chat-%d", chat.ID), redis_rate.PerSecond(1)); err != nil { + otelzap.L().Ctx(ctx).Error("failed to rate limit", zap.Error(err)) + return + } else if res.Allowed == 0 { + otelzap.L().Ctx(ctx).Warn("rate limit exceeded") + return + } + + if _, err := bot.SendMessage( + telegoutil. + Message(chat.ChatID(), `DC1: 美国 迈阿密 +DC2: 荷兰 阿姆斯特丹 +DC3: 美国 迈阿密 +DC4: 荷兰 阿姆斯特丹 +DC5: 新加坡 + +注册手机区号对应数据中心信息`). + WithParseMode(telego.ModeHTML), + ); err != nil { + otelzap.L().Ctx(ctx).Error("failed to send message", zap.Error(err)) + return + } +} diff --git a/components/basic-handler/internal/bot/handlers/dc.tpl b/components/basic-handler/internal/bot/handlers/dc.tpl new file mode 100644 index 0000000..2c33b7c --- /dev/null +++ b/components/basic-handler/internal/bot/handlers/dc.tpl @@ -0,0 +1,13 @@ +{{ if ne .Message.Chat.Username "" -}} +此群组所在数据中心为 DC{{ datacenter .Message.Chat.Username }} +{{ else -}} +此群组未设置用户名 +{{- end }} + +{{ if and .Message.From (ne .Message.From.Username "") -}} +您所在数据中心为 DC{{ datacenter .Message.From.Username }} +{{ else -}} +您未设置用户名 +{{- end }} + +此数据中心数据通过聊天头像查询，不保证准确性。 \ No newline at end of file diff --git a/components/basic-handler/internal/bot/handlers/handlers.go b/components/basic-handler/internal/bot/handlers/handlers.go index 3467481..f10855a 100644 --- a/components/basic-handler/internal/bot/handlers/handlers.go +++ b/components/basic-handler/internal/bot/handlers/handlers.go @@ -1,11 +1,13 @@ package handlers import ( + "github.com/go-redis/redis_rate/v10" "github.com/mymmrac/telego" "go.opentelemetry.io/otel" "go.uber.org/fx" "github.com/ahdark-services/pegasus/components/basic-handler/services/action_reply" + "github.com/ahdark-services/pegasus/components/basic-handler/services/datacenter" ) var tracer = otel.Tracer("github.com/ahdark-services/pegasus/components/remake-handler/internal/bot/handlers") @@ -14,11 +16,15 @@ type Handlers interface { StartCommandHandler(bot *telego.Bot, update telego.Update) ActionCommandHandler(bot *telego.Bot, update telego.Update) IDCommandHandler(bot *telego.Bot, update telego.Update) + DatacenterCommandHandler(bot *telego.Bot, update telego.Update) + DatacenterMoreInfoHandler(bot *telego.Bot, update telego.Update) } type handlers struct { fx.In ActionReplyService action_reply.Service + DatacenterService datacenter.Service + RedisRateLimiter *redis_rate.Limiter } func NewHandlers(h handlers) Handlers { diff --git a/components/basic-handler/services/0module.go b/components/basic-handler/services/0module.go index a79ad7b..1b573d8 100644 --- a/components/basic-handler/services/0module.go +++ b/components/basic-handler/services/0module.go @@ -1,12 +1,15 @@ package services import ( - "github.com/ahdark-services/pegasus/components/basic-handler/services/action_reply" "go.uber.org/fx" + + "github.com/ahdark-services/pegasus/components/basic-handler/services/action_reply" + "github.com/ahdark-services/pegasus/components/basic-handler/services/datacenter" ) func Module() fx.Option { return fx.Module("services", fx.Provide(action_reply.NewService), + fx.Provide(datacenter.NewService), ) } diff --git a/components/basic-handler/services/datacenter/service.go b/components/basic-handler/services/datacenter/service.go new file mode 100644 index 0000000..60e159a --- /dev/null +++ b/components/basic-handler/services/datacenter/service.go @@ -0,0 +1,79 @@ +package datacenter + +import ( + "context" + "regexp" + + "github.com/ahdark-services/pegasus/pkg/utils" + "github.com/bytedance/sonic" + "github.com/imroc/req/v3" + "github.com/uptrace/opentelemetry-go-extra/otelzap" + "go.opentelemetry.io/otel" + "go.uber.org/zap" +) + +var tracer = otel.Tracer("github.com/ahdark-services/pegasus/components/gateway/services/datacenter") + +type Service interface { + QueryDatacenterByUsername(ctx context.Context, username string) (int, error) +} + +type service struct { + client *req.Client +} + +func NewService() Service { + client := req.NewClient(). + SetJsonMarshal(sonic.Marshal). + SetJsonUnmarshal(sonic.Unmarshal). + SetBaseURL("https://t.me"). + SetLogger(otelzap.L().Named("service.datacenter").Sugar()). + SetCommonRetryCount(3). + SetCommonRetryCondition(func(resp *req.Response, err error) bool { + return resp.Response.StatusCode >= 500 + }). + SetCommonRetryHook(func(resp *req.Response, err error) { + otelzap.L().Ctx(resp.Request.Context()). + Error("failed to do request", + zap.Error(err), + zap.Int("response.status_code", resp.Response.StatusCode), + zap.String("request.method", resp.Request.Method), + zap.String("request.url", resp.Request.URL.String()), + ) + }). + EnableDumpEachRequest(). + WrapRoundTripFunc(utils.TraceRoundTripWrapperFunc(tracer, "DatacenterService.client.RoundTrip")) + + return &service{client} +} + +var dcRegexp = regexp.MustCompile(`https://cdn(\d).cdn-telegram.org/file/[\w-_]+\.\w+`) + +func (svc *service) QueryDatacenterByUsername(ctx context.Context, username string) (int, error) { + ctx, span := tracer.Start(ctx, "DatacenterService.QueryDatacenterByUsername") + defer span.End() + + resp, err := svc.client.R().SetPathParam("username", username).Get("/{username}") + if err != nil { + otelzap.L().Ctx(ctx).Error("failed to do request", zap.Error(err)) + return 0, err + } + + if !resp.IsSuccessState() { + otelzap.L().Ctx(ctx).Error("failed to do request", zap.String("response.body", resp.String())) + return 0, err + } + + bodyContent, err := resp.ToString() + if err != nil { + otelzap.L().Ctx(ctx).Error("failed to do request", zap.Error(err)) + return 0, err + } + + matches := dcRegexp.FindStringSubmatch(bodyContent) + if len(matches) < 2 { + return 0, nil + } + + return int(matches[1][0] - '0'), nil +} diff --git a/components/basic-handler/services/datacenter/service_test.go b/components/basic-handler/services/datacenter/service_test.go new file mode 100644 index 0000000..cd0241b --- /dev/null +++ b/components/basic-handler/services/datacenter/service_test.go @@ -0,0 +1,22 @@ +package datacenter + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestNewService(t *testing.T) { + asserts := assert.New(t) + asserts.NotNil(NewService()) +} + +func TestService_GetDatacenter(t *testing.T) { + asserts := assert.New(t) + svc := NewService() + + dc, err := svc.QueryDatacenterByUsername(context.Background(), "durov") + asserts.NoError(err) + asserts.Equal(1, dc) +}