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