add ApiKey limit

AAStarCommunity · Jun 20, 2024 · 484bdc8 · 484bdc8
1 parent f3c7197
commit 484bdc8
Show file tree

Hide file tree

Showing 8 changed files with 121 additions and 48 deletions.
diff --git a/common/model/api_key_info.go b/common/model/api_key_info.go
@@ -1,17 +1,20 @@
 package model
 
-import "golang.org/x/time/rate"
+import (
+	mapset "github.com/deckarep/golang-set/v2"
+	"golang.org/x/time/rate"
+)
 
 type ApiKeyModel struct {
-	Disable                       bool       `json:"disable"`
-	ApiKey                        string     `json:"api_key"`
-	RateLimit                     rate.Limit `json:"rate_limit"`
-	UserId                        int64      `json:"user_id"`
-	NetWorkLimitEnable            bool       `json:"network_limit_enable"`
-	DomainWhitelist               []string   `json:"domain_whitelist"`
-	IPWhiteList                   []string   `json:"ip_white_list"`
-	PaymasterEnable               bool       `json:"paymaster_enable"`
-	Erc20PaymasterEnable          bool       `json:"erc20_paymaster_enable"`
-	ProjectSponsorPaymasterEnable bool       `json:"project_sponsor_paymaster_enable"`
-	UserPayPaymasterEnable        bool       `json:"user_pay_paymaster_enable"`
+	Disable                       bool               `json:"disable"`
+	ApiKey                        string             `json:"api_key"`
+	RateLimit                     rate.Limit         `json:"rate_limit"`
+	UserId                        int64              `json:"user_id"`
+	NetWorkLimitEnable            bool               `json:"network_limit_enable"`
+	DomainWhitelist               mapset.Set[string] `json:"domain_whitelist"`
+	IPWhiteList                   mapset.Set[string] `json:"ip_white_list"`
+	PaymasterEnable               bool               `json:"paymaster_enable"`
+	Erc20PaymasterEnable          bool               `json:"erc20_paymaster_enable"`
+	ProjectSponsorPaymasterEnable bool               `json:"project_sponsor_paymaster_enable"`
+	UserPayPaymasterEnable        bool               `json:"user_pay_paymaster_enable"`
 }
diff --git a/rpc_server/middlewares/auth.go b/rpc_server/middlewares/auth.go
@@ -3,36 +3,54 @@ package middlewares
 import (
 	"AAStarCommunity/EthPaymaster_BackService/common/global_const"
 	"AAStarCommunity/EthPaymaster_BackService/service/dashboard_service"
+	"errors"
 	"github.com/gin-gonic/gin"
-	"github.com/sirupsen/logrus"
 	"net/http"
 )
 
 type ApiKey struct {
 	Key string `form:"apiKey" json:"apiKey" binding:"required"`
 }
 
-func AuthHandler() gin.HandlerFunc {
+func ApiVerificationHandler() gin.HandlerFunc {
 	return func(c *gin.Context) {
 		apiKey := c.Query("apiKey")
 		if apiKey == "" {
-			c.JSON(http.StatusForbidden, gin.H{"error": "ApiKey is mandatory, visit to https://dashboard.aastar.io for more detail."})
-			c.Abort()
+			_ = c.AbortWithError(http.StatusForbidden, errors.New("ApiKey is mandatory, visit to https://dashboard.aastar.io for more detail"))
 			return
 		}
 		apiModel, err := dashboard_service.GetAPiInfoByApiKey(apiKey)
 		if err != nil {
-			logrus.Errorf("GetAPiInfoByApiKey err: %v", err)
-			c.JSON(http.StatusBadRequest, gin.H{"error": "Can Not Find Your Api Key"})
-			c.Abort()
+			_ = c.AbortWithError(http.StatusBadRequest, errors.New("can Not Find Your Api Key"))
 			return
 		}
 		if apiModel.Disable {
-			c.JSON(http.StatusForbidden, gin.H{"error": "Api Key Is Disabled"})
-			c.Abort()
+			_ = c.AbortWithError(http.StatusForbidden, errors.New("api Key Is Disabled"))
+			return
+		}
+		if !apiModel.PaymasterEnable {
+			_ = c.AbortWithError(http.StatusForbidden, errors.New("api Key Is Disabled Paymaster"))
+			return
+		}
+		if !VerifyRateLimit(*apiModel) {
+			_ = c.AbortWithError(http.StatusTooManyRequests, errors.New("too many requests"))
 			return
 		}
 
+		if apiModel.IPWhiteList != nil && apiModel.IPWhiteList.Cardinality() > 0 {
+			clientIp := c.ClientIP()
+			if !apiModel.IPWhiteList.Contains(clientIp) {
+				_ = c.AbortWithError(http.StatusForbidden, errors.New("ip not in whitelist"))
+				return
+			}
+		}
+		if apiModel.DomainWhitelist != nil && apiModel.DomainWhitelist.Cardinality() > 0 {
+			domain := c.Request.Host
+			if !apiModel.DomainWhitelist.Contains(domain) {
+				_ = c.AbortWithError(http.StatusForbidden, errors.New("domain not in whitelist"))
+				return
+			}
+		}
 		c.Set(global_const.ContextKeyApiMoDel, apiModel)
 	}
 }
diff --git a/rpc_server/middlewares/pv_mertics.go b/rpc_server/middlewares/pv_mertics.go
@@ -69,7 +69,14 @@ func PvMetrics() gin.HandlerFunc {
 				metricsParam.ApiKey = apiKeyModel.ApiKey
 				metricsParam.ApiUserId = apiKeyModel.UserId
 			}
-			metricsPaymaster(c, metricsParam)
+
+			extraMap := make(map[string]any)
+			clientIp := c.ClientIP()
+			extraMap["client_ip"] = clientIp
+
+			domain := c.Request.Host
+			extraMap["client_domain"] = domain
+			metricsPaymaster(c, metricsParam, &extraMap)
 		} else {
 			return
 		}
@@ -86,7 +93,7 @@ func (w *CustomResponseWriter) Write(b []byte) (int, error) {
 	w.body.Write(b)
 	return w.ResponseWriter.Write(b)
 }
-func metricsPaymaster(c *gin.Context, metricsParam PayMasterParam) {
+func metricsPaymaster(c *gin.Context, metricsParam PayMasterParam, extraMap *map[string]any) {
 
 	recallModel := dashboard_service.PaymasterRecallLogDbModel{
 		ProjectApikey:   metricsParam.ApiKey,
@@ -99,6 +106,13 @@ func metricsPaymaster(c *gin.Context, metricsParam PayMasterParam) {
 		Status:          metricsParam.Status,
 		NetWork:         metricsParam.NetWork,
 	}
+	if extraMap != nil {
+		executeRestrictionJson, err := json.Marshal(extraMap)
+		if err != nil {
+			logrus.Error("executeRestrictionJson error:", err)
+		}
+		recallModel.Extra = executeRestrictionJson
+	}
 	err := dashboard_service.CreatePaymasterCall(&recallModel)
 	if err != nil {
 		logrus.Error("CreatePaymasterCall error:", err)

diff --git a/rpc_server/middlewares/rate_limit.go b/rpc_server/middlewares/rate_limit.go
@@ -1,12 +1,8 @@
 package middlewares
 
 import (
-	"AAStarCommunity/EthPaymaster_BackService/common/global_const"
 	"AAStarCommunity/EthPaymaster_BackService/common/model"
-	"errors"
-	"github.com/gin-gonic/gin"
 	"golang.org/x/time/rate"
-	"net/http"
 )
 
 const (
@@ -16,21 +12,10 @@ const (
 
 var limiter map[string]*rate.Limiter
 
-// RateLimiterByApiKeyHandler represents the rate limit by each ApiKey for each api calling
-func RateLimiterByApiKeyHandler() gin.HandlerFunc {
-	return func(ctx *gin.Context) {
-		apiKeyModelInterface := ctx.MustGet(global_const.ContextKeyApiMoDel)
-		defaultLimit := DefaultLimit
-		apiKeyModel := apiKeyModelInterface.(*model.ApiKeyModel)
-		defaultLimit = apiKeyModel.RateLimit
-
-		if limiting(&apiKeyModel.ApiKey, defaultLimit) {
-			ctx.Next()
-		} else {
-			_ = ctx.AbortWithError(http.StatusTooManyRequests, errors.New("too many requests"))
-		}
-	}
+func VerifyRateLimit(keyModel model.ApiKeyModel) bool {
+	return limiting(&keyModel.ApiKey, keyModel.RateLimit)
 }
+
 func clearLimiter(apiKey *string) {
 	delete(limiter, *apiKey)
 }

diff --git a/rpc_server/routers/boot.go b/rpc_server/routers/boot.go
@@ -52,8 +52,7 @@ func buildRoute(routers *gin.Engine) {
 	//build the routers not need api access like auth or Traffic limit
 	buildRouters(routers, PublicRouterMaps)
 
-	routers.Use(middlewares.AuthHandler())
-	routers.Use(middlewares.RateLimiterByApiKeyHandler())
+	routers.Use(middlewares.ApiVerificationHandler())
 	buildRouters(routers, PrivateRouterMaps)
 }
 

diff --git a/service/dashboard_service/dashboard_service.go b/service/dashboard_service/dashboard_service.go
@@ -234,13 +234,52 @@ func (*ApiKeyDbModel) TableName() string {
 func (m *ApiKeyDbModel) GetRateLimit() rate.Limit {
 	return 10
 }
+
+type APiModelExtra struct {
+	NetWorkLimitEnable            bool     `json:"network_limit_enable"`
+	DomainWhitelist               []string `json:"domain_whitelist"`
+	IPWhiteList                   []string `json:"ip_white_list"`
+	PaymasterEnable               bool     `json:"paymaster_enable"`
+	Erc20PaymasterEnable          bool     `json:"erc20_paymaster_enable"`
+	ProjectSponsorPaymasterEnable bool     `json:"project_sponsor_paymaster_enable"`
+	UserPayPaymasterEnable        bool     `json:"user_pay_paymaster_enable"`
+}
+
 func convertApiKeyDbModelToApiKeyModel(apiKeyDbModel *ApiKeyDbModel) *model.ApiKeyModel {
-	return &model.ApiKeyModel{
+	apiKeyModel := &model.ApiKeyModel{
 		Disable:   apiKeyDbModel.Disable,
 		ApiKey:    apiKeyDbModel.ApiKey,
 		RateLimit: 10,
 		UserId:    apiKeyDbModel.UserId,
 	}
+	if apiKeyDbModel.Extra != nil {
+		// convert To map
+		eJson, _ := apiKeyDbModel.Extra.MarshalJSON()
+		apiKeyExtra := &APiModelExtra{}
+		err := json.Unmarshal(eJson, apiKeyExtra)
+		if err != nil {
+			return nil
+		}
+		apiKeyModel.NetWorkLimitEnable = apiKeyExtra.NetWorkLimitEnable
+		if apiKeyExtra.IPWhiteList != nil {
+			apiKeyModel.IPWhiteList = mapset.NewSetWithSize[string](len(apiKeyExtra.IPWhiteList))
+			for _, v := range apiKeyExtra.IPWhiteList {
+				apiKeyModel.IPWhiteList.Add(v)
+			}
+		}
+		if apiKeyExtra.DomainWhitelist != nil {
+			apiKeyModel.DomainWhitelist = mapset.NewSetWithSize[string](len(apiKeyExtra.DomainWhitelist))
+			for _, v := range apiKeyExtra.DomainWhitelist {
+				apiKeyModel.DomainWhitelist.Add(v)
+			}
+		}
+		apiKeyModel.PaymasterEnable = apiKeyExtra.PaymasterEnable
+		apiKeyModel.Erc20PaymasterEnable = apiKeyExtra.Erc20PaymasterEnable
+		apiKeyModel.ProjectSponsorPaymasterEnable = apiKeyExtra.ProjectSponsorPaymasterEnable
+		apiKeyModel.UserPayPaymasterEnable = apiKeyExtra.UserPayPaymasterEnable
+
+	}
+	return apiKeyModel
 }
 func GetAPiInfoByApiKey(apiKey string) (*model.ApiKeyModel, error) {
 	apikeyModel := &ApiKeyDbModel{}

diff --git a/service/operator/try_pay_user_op_execute.go b/service/operator/try_pay_user_op_execute.go
@@ -20,7 +20,7 @@ import (
 )
 
 func TryPayUserOpExecute(apiKeyModel *model.ApiKeyModel, request *model.UserOpRequest) (*model.TryPayUserOpResponse, error) {
-	userOp, strategy, paymasterDataInput, err := prepareExecute(request)
+	userOp, strategy, paymasterDataInput, err := prepareExecute(request, apiKeyModel)
 	if err != nil {
 		return nil, err
 	}
@@ -46,14 +46,14 @@ func TryPayUserOpExecute(apiKeyModel *model.ApiKeyModel, request *model.UserOpRe
 	return result, nil
 }
 
-func prepareExecute(request *model.UserOpRequest) (*user_op.UserOpInput, *model.Strategy, *paymaster_data.PaymasterDataInput, error) {
+func prepareExecute(request *model.UserOpRequest, apiKeyModel *model.ApiKeyModel) (*user_op.UserOpInput, *model.Strategy, *paymaster_data.PaymasterDataInput, error) {
 	var strategy *model.Strategy
 	strategy, generateErr := StrategyGenerate(request)
 	if generateErr != nil {
 		return nil, nil, nil, generateErr
 	}
 
-	if err := validator_service.ValidateStrategy(strategy, request); err != nil {
+	if err := validator_service.ValidateStrategy(strategy, request, apiKeyModel); err != nil {
 		return nil, nil, nil, err
 	}
 

diff --git a/service/validator_service/basic_validator.go b/service/validator_service/basic_validator.go
@@ -11,7 +11,7 @@ import (
 	"time"
 )
 
-func ValidateStrategy(strategy *model.Strategy, request *model.UserOpRequest) error {
+func ValidateStrategy(strategy *model.Strategy, request *model.UserOpRequest, keyModel *model.ApiKeyModel) error {
 	if strategy == nil {
 		return xerrors.Errorf("empty strategy")
 	}
@@ -77,6 +77,21 @@ func ValidateStrategy(strategy *model.Strategy, request *model.UserOpRequest) er
 			return xerrors.Errorf("strategy not support chainId [%s]", netWorkStr)
 		}
 	}
+	payType := strategy.GetPayType()
+	switch payType {
+	case global_const.PayTypeERC20:
+		if !keyModel.Erc20PaymasterEnable {
+			return xerrors.Errorf("strategy pay type is erc20 but not enable")
+		}
+	case global_const.PayTypeUserSponsor:
+		if !keyModel.UserPayPaymasterEnable {
+			return xerrors.Errorf("strategy pay type is user sponsor but not enable")
+		}
+	case global_const.PayTypeProjectSponsor:
+		if !keyModel.ProjectSponsorPaymasterEnable {
+			return xerrors.Errorf("strategy pay type is project sponsor but not enable")
+		}
+	}
 
 	return nil