test: 添加tma测试代码

TBXark · Oct 24, 2024 · 4deecc5 · 4deecc5
1 parent 3f27d95
commit 4deecc5
Show file tree

Hide file tree

Showing 4 changed files with 100 additions and 24 deletions.
diff --git a/go.mod b/go.mod
@@ -12,19 +12,19 @@ require (
 	github.com/gin-contrib/zap v1.1.4
 	github.com/gin-gonic/gin v1.10.0
 	github.com/go-sql-driver/mysql v1.8.1
-	github.com/go-telegram/bot v1.9.0
+	github.com/go-telegram/bot v1.9.1
 	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/google/wire v0.6.0
 	github.com/meilisearch/meilisearch-go v0.28.0
 	github.com/minio/minio-go/v7 v7.0.78
 	github.com/patrickmn/go-cache v2.1.0+incompatible
-	github.com/qiniu/go-sdk/v7 v7.23.0
+	github.com/qiniu/go-sdk/v7 v7.24.0
 	github.com/redis/go-redis/v9 v9.7.0
 	github.com/samber/lo v1.47.0
 	github.com/spf13/cobra v1.8.1
 	github.com/swaggo/files v1.0.1
 	github.com/swaggo/gin-swagger v1.6.0
-	github.com/swaggo/swag v1.16.3
+	github.com/swaggo/swag v1.16.4
 	github.com/telegram-mini-apps/init-data-golang v1.1.5
 	github.com/yitter/idgenerator-go v1.3.3
 	go.uber.org/zap v1.27.0

diff --git a/go.sum b/go.sum
@@ -105,6 +105,8 @@ github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpv
 github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
 github.com/go-telegram/bot v1.9.0 h1:z9g0Fgk9B7G/xoVMqji30hpJPlr3Dz3aVW2nzSGfPuI=
 github.com/go-telegram/bot v1.9.0/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
+github.com/go-telegram/bot v1.9.1 h1:4vkNV6vDmEPZaYP7sZYaagOaJyV4GerfOPkjg/Ki5ic=
+github.com/go-telegram/bot v1.9.1/go.mod h1:i2TRs7fXWIeaceF3z7KzsMt/he0TwkVC680mvdTFYeM=
 github.com/go-test/deep v1.0.3 h1:ZrJSEWsXzPOxaZnFteGEfooLba+ju3FYIbOrS+rQd68=
 github.com/go-test/deep v1.0.3/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
 github.com/goccy/go-json v0.10.3 h1:KZ5WoDbxAIgm2HNbYckL0se1fHD6rz5j4ywS6ebzDqA=
@@ -201,6 +203,8 @@ github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoG
 github.com/qiniu/dyn v1.3.0/go.mod h1:E8oERcm8TtwJiZvkQPbcAh0RL8jO1G0VXJMW3FAWdkk=
 github.com/qiniu/go-sdk/v7 v7.23.0 h1:4wYB4EGE6MBhvjtE/FZH/mIUt/VH6WjzBucU3VfPwhg=
 github.com/qiniu/go-sdk/v7 v7.23.0/go.mod h1:OXsAVU5YrLLtVi4iPFpP80jzb3SRBAczrGkcqQmWhcY=
+github.com/qiniu/go-sdk/v7 v7.24.0 h1:1Vq8Xb1GrwqpeLp1HAnKv6xi0d9GyyiAapJb5/X2na8=
+github.com/qiniu/go-sdk/v7 v7.24.0/go.mod h1:uZE85Pi0ftIHT/UNLShosdzwsovqpdas0LwAGO7cPao=
 github.com/qiniu/x v1.10.5/go.mod h1:03Ni9tj+N2h2aKnAz+6N0Xfl8FwMEDRC2PAlxekASDs=
 github.com/redis/go-redis/v9 v9.7.0 h1:HhLSs+B6O021gwzl+locl0zEDnyNkxMtf/Z3NNBMa9E=
 github.com/redis/go-redis/v9 v9.7.0/go.mod h1:f6zhXITC7JUJIlPEiBOTXxJgPLdZcA93GewI7inzyWw=
@@ -240,6 +244,8 @@ github.com/swaggo/gin-swagger v1.6.0 h1:y8sxvQ3E20/RCyrXeFfg60r6H0Z+SwpTjMYsMm+z
 github.com/swaggo/gin-swagger v1.6.0/go.mod h1:BG00cCEy294xtVpyIAHG6+e2Qzj/xKlRdOqDkvq0uzo=
 github.com/swaggo/swag v1.16.3 h1:PnCYjPCah8FK4I26l2F/KQ4yz3sILcVUN3cTlBFA9Pg=
 github.com/swaggo/swag v1.16.3/go.mod h1:DImHIuOFXKpMFAQjcC7FG4m3Dg4+QuUgUzJmKjI/gRk=
+github.com/swaggo/swag v1.16.4 h1:clWJtd9LStiG3VeijiCfOVODP6VpHtKdQy9ELFG3s1A=
+github.com/swaggo/swag v1.16.4/go.mod h1:VBsHJRsDvfYvqoiMKnsdwhNV9LEMHgEDZcyVYX0sxPg=
 github.com/telegram-mini-apps/init-data-golang v1.1.5 h1:R51eoGSKBQwHoAo8r/n/E0RZ2owF3kmEpdzn7oV7lgI=
 github.com/telegram-mini-apps/init-data-golang v1.1.5/go.mod h1:GG4HnRx9ocjD4MjjzOw7gf9Ptm0NvFbDr5xqnfFOYuY=
 github.com/twitchyliquid64/golang-asm v0.15.1 h1:SU5vSMR7hnwNxj24w34ZyCi/FmDZTkS4MhqMhdFk5YI=
@@ -367,6 +373,7 @@ modernc.org/cc/v4 v4.21.4 h1:3Be/Rdo1fpr8GrQ7IVw9OHtplU4gWbb+wNgeoBMmGLQ=
 modernc.org/cc/v4 v4.21.4/go.mod h1:HM7VJTZbUCR3rV8EYBi9wxnJ0ZBRiGE5OeGXNA0IsLQ=
 modernc.org/ccgo/v4 v4.21.0 h1:kKPI3dF7RIag8YcToh5ZwDcVMIv6VGa0ED5cvh0LMW4=
 modernc.org/ccgo/v4 v4.21.0/go.mod h1:h6kt6H/A2+ew/3MW/p6KEoQmrq/i3pr0J/SiwiaF/g0=
+modernc.org/fileutil v1.0.0/go.mod h1:JHsWpkrk/CnVV1H/eGlFf85BEpfkrp56ro8nojIq9Q8=
 modernc.org/fileutil v1.3.0 h1:gQ5SIzK3H9kdfai/5x41oQiKValumqNTDXMvKo62HvE=
 modernc.org/fileutil v1.3.0/go.mod h1:XatxS8fZi3pS8/hKG2GH/ArUogfxjpEKs3Ku3aK4JyQ=
 modernc.org/gc/v2 v2.5.0 h1:bJ9ChznK1L1mUtAQtxi0wi5AtAs5jQuw4PrPHO5pb6M=

diff --git a/pkg/server/auth/tmaauth/tma.go b/pkg/server/auth/tmaauth/tma.go
@@ -1,35 +1,27 @@
 package tmaauth
 
 import (
-	"github.com/tbxark/sphere/pkg/server/auth/authorizer"
+	"bytes"
+	"encoding/json"
 	initdata "github.com/telegram-mini-apps/init-data-golang"
+	"net/url"
+	"strconv"
 	"time"
 )
 
 const AuthorizationPrefixTMA = "tma"
 
-type Claims struct {
-	UID       int64              `json:"uid"`
-	InitData  *initdata.InitData `json:"init_data"`
-	ExpiresAt time.Time          `json:"exp"`
-}
-
-func (c *Claims) Valid() error {
-	if c.ExpiresAt.Before(time.Now()) {
-		return authorizer.ErrorExpiredToken
-	}
-	return nil
-}
+type Claims initdata.InitData
 
 type TmaAuth struct {
 	token string
 	expIn time.Duration
 }
 
-func NewTmaAuth(token string) *TmaAuth {
+func NewTmaAuth(token string, expIn time.Duration) *TmaAuth {
 	return &TmaAuth{
 		token: token,
-		expIn: time.Hour * 24,
+		expIn: expIn,
 	}
 }
 
@@ -38,14 +30,50 @@ func (t *TmaAuth) ParseToken(token string) (*Claims, error) {
 	if err != nil {
 		return nil, err
 	}
-	initData, err := initdata.Parse(token)
+	data, err := initdata.Parse(token)
 	if err != nil {
 		return nil, err
 	}
-	claims := Claims{
-		UID:       initData.User.ID,
-		InitData:  &initData,
-		ExpiresAt: time.Now().Add(t.expIn),
+	return (*Claims)(&data), nil
+}
+
+func (t *TmaAuth) GenerateToken(claims *Claims) (string, error) {
+	rawInitMap := map[string]any{}
+	initBytes, err := json.Marshal(claims)
+	if err != nil {
+		return "", err
+	}
+	decoder := json.NewDecoder(bytes.NewReader(initBytes))
+	decoder.UseNumber()
+	err = decoder.Decode(&rawInitMap)
+	if err != nil {
+		return "", err
+	}
+	delete(rawInitMap, "hash")
+	delete(rawInitMap, "auth_date")
+	params := make(map[string]string)
+	values := url.Values{}
+	for k, v := range rawInitMap {
+		if str, ok := v.(string); ok {
+			params[k] = v.(string)
+			values.Add(k, str)
+			continue
+		}
+		if num, ok := v.(json.Number); ok {
+			params[k] = num.String()
+			values.Add(k, num.String())
+			continue
+		}
+		partBytes, e := json.Marshal(v)
+		if e != nil {
+			return "", e
+		}
+		params[k] = string(partBytes)
+		values.Add(k, string(partBytes))
 	}
-	return &claims, nil
+	exp := time.Unix(int64(claims.AuthDateRaw), 0)
+	sign := initdata.Sign(params, t.token, exp)
+	values.Set("hash", sign)
+	values.Set("auth_date", strconv.FormatInt(exp.Unix(), 10))
+	return values.Encode(), nil
 }
diff --git a/pkg/server/auth/tmaauth/tma_test.go b/pkg/server/auth/tmaauth/tma_test.go
@@ -0,0 +1,41 @@
+package tmaauth
+
+import (
+	initdata "github.com/telegram-mini-apps/init-data-golang"
+	"testing"
+	"time"
+)
+
+func TestTmaAuth_ParseToken(t *testing.T) {
+	secretToken := "test"
+	claims := Claims{
+		ChatInstance:    123,
+		CanSendAfterRaw: 1234,
+		User: initdata.User{
+			AddedToAttachmentMenu: false,
+			AllowsWriteToPm:       false,
+			FirstName:             "full_name",
+			ID:                    1,
+			IsBot:                 false,
+			IsPremium:             false,
+			LastName:              "last_name",
+			Username:              "username",
+			LanguageCode:          "",
+			PhotoURL:              "",
+		},
+		AuthDateRaw: int(time.Now().Unix()),
+	}
+	tmaAuth := NewTmaAuth(secretToken, time.Hour)
+	token, err := tmaAuth.GenerateToken(&claims)
+	if err != nil {
+		t.Error(err)
+	}
+	t.Log(token)
+	parsedClaims, err := tmaAuth.ParseToken(token)
+	if err != nil {
+		t.Error(err)
+	}
+	if parsedClaims.User.ID != claims.User.ID {
+		t.Errorf("expected %d, got %d", claims.User.ID, parsedClaims.User.ID)
+	}
+}