Merge pull request #4 from ahdark-services/feat/datacenter_query

Allow users to query their data center

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()))
 }

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: 新加坡
+
+<a href="https://t.me/KinhRoBotChannel/88">注册手机区号对应数据中心信息</a>`).
+ WithParseMode(telego.ModeHTML),
+ ); err != nil {
+ otelzap.L().Ctx(ctx).Error("failed to send message", zap.Error(err))
+ return
+ }
+}

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 }}
+
+此数据中心数据通过聊天头像查询，不保证准确性。

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 {

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),
  )
 }

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
+}

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)
+}