Skip to content

Commit 6538d48

Browse files
Dev AgentQinYuuuu
Dev Agent
authored and
QinYuuuu
committed
Feat api limit for ip geo
1 parent cf4d0d2 commit 6538d48

File tree

2 files changed

+61
-19
lines changed

2 files changed

+61
-19
lines changed

‎api/router/api.go‎

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,10 @@ func NewRouter(config *config.Config, enableSwagger bool) (*gin.Engine, error) {
8181
middlewareCollection.Auth.NeedPhoneVerified = middleware.NeedPhoneVerified(config)
8282
middlewareCollection.Repo.RepoExists = middleware.RepoExists(config)
8383
middlewareCollection.License.Check = middleware.CheckLicense(config)
84-
84+
middlewareCollection.API.Captcha = middleware.NewCaptcha(config).Handle
85+
rateLimitMiddleware := middleware.RateLimiter(config, middleware.WithTimeBucketRateLimter(config), middleware.WithIPCheck())
86+
middlewareCollection.API.IPLimiter = middleware.IPLocationCheck(config)
87+
middlewareCollection.API.RateLimter = rateLimitMiddleware
8588
//add router for golang pprof
8689
debugGroup := r.Group("/debug", middlewareCollection.Auth.NeedAPIKey)
8790
pprof.RouteRegister(debugGroup, "pprof")
@@ -183,14 +186,12 @@ func NewRouter(config *config.Config, enableSwagger bool) (*gin.Engine, error) {
183186

184187
r.Use(middleware.LocalizedErrorMiddleware())
185188
r.Use(middleware.Authenticator(config))
189+
r.Use(middlewareCollection.API.IPLimiter, middlewareCollection.API.RateLimter, middlewareCollection.API.Captcha)
186190
apiGroup := r.Group("/api/v1")
187191

188192
versionHandler := handler.NewVersionHandler()
189193
apiGroup.GET("/version", versionHandler.Version)
190194

191-
// Admin user get repo path list
192-
apiGroup.GET("/repos", middlewareCollection.Auth.NeedAdmin, repoCommonHandler.GetRepos)
193-
194195
// TODO:use middleware to handle common response
195196
//
196197
memoryStore := persist.NewMemoryStore(1 * time.Minute)
@@ -314,6 +315,7 @@ func NewRouter(config *config.Config, enableSwagger bool) (*gin.Engine, error) {
314315
apiGroup.POST("/jwt/token", middlewareCollection.Auth.NeedAPIKey, userProxyHandler.Proxy)
315316
apiGroup.GET("/jwt/:token", middlewareCollection.Auth.NeedAPIKey, userProxyHandler.ProxyToApi("/api/v1/jwt/%s", "token"))
316317
apiGroup.GET("/users", userProxyHandler.Proxy)
318+
apiGroup.GET("/users/stream-export", middlewareCollection.Auth.NeedAdmin, userProxyHandler.Proxy)
317319

318320
// callback
319321
callbackCtrl, err := callback.NewGitCallbackHandler(config)
@@ -404,6 +406,9 @@ func NewRouter(config *config.Config, enableSwagger bool) (*gin.Engine, error) {
404406

405407
adminGroup.POST("/:repo_type/:namespace/:name/change_path", repoCommonHandler.ChangePath)
406408

409+
// Admin user get repo path list
410+
adminGroup.GET("/repos", repoCommonHandler.GetRepos)
411+
407412
// routes for broadcast
408413
broadcastHandler, err := handler.NewBroadcastHandler()
409414
if err != nil {
@@ -554,6 +559,12 @@ func NewRouter(config *config.Config, enableSwagger bool) (*gin.Engine, error) {
554559
return nil, fmt.Errorf("error creating webhook routes: %w", err)
555560
}
556561

562+
// Initialize LFS component for Xnet processing
563+
_, err = handler.NewLfsHandler(config, mqFactory)
564+
if err != nil {
565+
return nil, fmt.Errorf("error creating lfs handler: %w", err)
566+
}
567+
557568
finetuneJobHandler, err := handler.NewFinetuneHandler(config)
558569
if err != nil {
559570
return nil, fmt.Errorf("error creating finetune job handler: %w", err)
@@ -601,7 +612,7 @@ func createModelRoutes(config *config.Config,
601612
modelsGroup.GET("", cache.Cache(memoryStore, time.Minute, middleware.CacheStrategyTrendingRepos()), modelHandler.Index)
602613
modelsGroup.PUT("/:namespace/:name", middlewareCollection.Auth.NeedLogin, modelHandler.Update)
603614
modelsGroup.DELETE("/:namespace/:name", middlewareCollection.Auth.NeedLogin, modelHandler.Delete)
604-
modelsGroup.GET("/:namespace/:name", modelHandler.Show)
615+
modelsGroup.GET("/:namespace/:name", cache.Cache(memoryStore, time.Minute*2, middleware.CacheRepoInfo()), modelHandler.Show)
605616
modelsGroup.GET("/:namespace/:name/relations", modelHandler.Relations)
606617
modelsGroup.PUT("/:namespace/:name/relations", middlewareCollection.Auth.NeedAdmin, modelHandler.SetRelations)
607618
modelsGroup.POST("/:namespace/:name/relations/dataset", middlewareCollection.Auth.NeedPhoneVerified, modelHandler.AddDatasetRelation)
@@ -615,7 +626,7 @@ func createModelRoutes(config *config.Config,
615626
modelsGroup.GET("/:namespace/:name/tags", repoCommonHandler.Tags)
616627
modelsGroup.POST("/:namespace/:name/preupload/:revision", middlewareCollection.Auth.NeedPhoneVerified, repoCommonHandler.Preupload)
617628
// update tags of a certain category
618-
modelsGroup.GET("/:namespace/:name/all_files", repoCommonHandler.AllFiles)
629+
modelsGroup.GET("/:namespace/:name/all_files", cache.Cache(memoryStore, time.Minute*2, middleware.CacheRepoInfo()), repoCommonHandler.AllFiles)
619630
modelsGroup.POST("/:namespace/:name/tags/:category", middlewareCollection.Auth.NeedPhoneVerified, repoCommonHandler.UpdateTags)
620631
modelsGroup.GET("/:namespace/:name/last_commit", repoCommonHandler.LastCommit)
621632
modelsGroup.GET("/:namespace/:name/commit/:commit_id", repoCommonHandler.CommitWithDiff)
@@ -627,7 +638,7 @@ func createModelRoutes(config *config.Config,
627638
modelsGroup.GET("/:namespace/:name/refs/:ref/remote_tree/*path", repoCommonHandler.RemoteTree)
628639
modelsGroup.GET("/:namespace/:name/refs/:ref/logs_tree/*path", repoCommonHandler.LogsTree)
629640
modelsGroup.GET("/:namespace/:name/commits", repoCommonHandler.Commits)
630-
modelsGroup.GET("/:namespace/:name/raw/*file_path", repoCommonHandler.FileRaw)
641+
modelsGroup.GET("/:namespace/:name/raw/*file_path", cache.Cache(memoryStore, time.Minute*2, middleware.CacheRepoInfo()), repoCommonHandler.FileRaw)
631642
modelsGroup.GET("/:namespace/:name/blob/*file_path", repoCommonHandler.FileInfo)
632643
// The DownloadFile method differs from the SDKDownload interface in a few ways
633644

@@ -743,6 +754,8 @@ func createDatasetRoutes(
743754
dsHandler *handler.DatasetHandler,
744755
repoCommonHandler *handler.RepoHandler,
745756
) {
757+
// gin cache
758+
memoryStore := persist.NewMemoryStore(2 * time.Minute)
746759
datasetsGroup := apiGroup.Group("/datasets")
747760
// allow access without login
748761
datasetsGroup.GET("", dsHandler.Index)
@@ -752,13 +765,13 @@ func createDatasetRoutes(
752765
datasetsGroup.POST("", middlewareCollection.Auth.NeedPhoneVerified, dsHandler.Create)
753766
datasetsGroup.PUT("/:namespace/:name", middleware.MustLogin(), dsHandler.Update)
754767
datasetsGroup.DELETE("/:namespace/:name", middleware.MustLogin(), dsHandler.Delete)
755-
datasetsGroup.GET("/:namespace/:name", dsHandler.Show)
768+
datasetsGroup.GET("/:namespace/:name", cache.Cache(memoryStore, time.Minute*2, middleware.CacheRepoInfo()), dsHandler.Show)
756769
datasetsGroup.GET("/:namespace/:name/relations", middleware.MustLogin(), dsHandler.Relations)
757770
datasetsGroup.GET("/:namespace/:name/branches", middleware.MustLogin(), repoCommonHandler.Branches)
758771
datasetsGroup.GET("/:namespace/:name/tags", middleware.MustLogin(), repoCommonHandler.Tags)
759772
datasetsGroup.POST("/:namespace/:name/preupload/:revision", middlewareCollection.Auth.NeedPhoneVerified, repoCommonHandler.Preupload)
760773
// update tags of a certain category
761-
datasetsGroup.GET("/:namespace/:name/all_files", middleware.MustLogin(), repoCommonHandler.AllFiles)
774+
datasetsGroup.GET("/:namespace/:name/all_files", middleware.MustLogin(), cache.Cache(memoryStore, time.Minute*2, middleware.CacheRepoInfo()), repoCommonHandler.AllFiles)
762775
datasetsGroup.POST("/:namespace/:name/tags/:category", middleware.MustLogin(), repoCommonHandler.UpdateTags)
763776
datasetsGroup.GET("/:namespace/:name/last_commit", repoCommonHandler.LastCommit)
764777
datasetsGroup.GET("/:namespace/:name/commit/:commit_id", middleware.MustLogin(), repoCommonHandler.CommitWithDiff)
@@ -771,7 +784,7 @@ func createDatasetRoutes(
771784
datasetsGroup.GET("/:namespace/:name/refs/:ref/logs_tree/*path", middleware.MustLogin(), repoCommonHandler.LogsTree)
772785
datasetsGroup.GET("/:namespace/:name/commits", middleware.MustLogin(), repoCommonHandler.Commits)
773786
datasetsGroup.POST("/:namespace/:name/raw/*file_path", middlewareCollection.Auth.NeedPhoneVerified, repoCommonHandler.CreateFile)
774-
datasetsGroup.GET("/:namespace/:name/raw/*file_path", middleware.MustLogin(), repoCommonHandler.FileRaw)
787+
datasetsGroup.GET("/:namespace/:name/raw/*file_path", middleware.MustLogin(), cache.Cache(memoryStore, time.Minute*2, middleware.CacheRepoInfo()), repoCommonHandler.FileRaw)
775788
datasetsGroup.GET("/:namespace/:name/blob/*file_path", repoCommonHandler.FileInfo)
776789
datasetsGroup.GET("/:namespace/:name/download/*file_path", middleware.MustLogin(), repoCommonHandler.DownloadFile)
777790
datasetsGroup.GET("/:namespace/:name/resolve/*file_path", middleware.MustLogin(), repoCommonHandler.ResolveDownload)
@@ -1122,6 +1135,8 @@ func createMappingRoutes(
11221135
hfDSAPIGroup.POST("/:namespace/:name/paths-info/:ref", middleware.RepoMapping(types.DatasetRepo), hfdsHandler.DatasetPathsInfo)
11231136
hfDSAPIGroup.GET("/:namespace/:name/tree/:ref/*path_in_repo", middleware.RepoMapping(types.DatasetRepo), hfdsHandler.DatasetTree)
11241137
hfDSAPIGroup.GET("/:namespace/:name/resolve/:ref/.huggingface.yaml", middleware.RepoMapping(types.DatasetRepo), hfdsHandler.HandleHFYaml)
1138+
hfDSAPIGroup.POST("/:namespace/:name/preupload/:revision", middleware.RepoMapping(types.DatasetRepo), repoCommonHandler.PreuploadHF)
1139+
hfDSAPIGroup.POST("/:namespace/:name/commit/:revision", middleware.RepoMapping(types.DatasetRepo), repoCommonHandler.CommitFilesHF)
11251140
}
11261141
hfSpaceAPIGroup := hfAPIGroup.Group("/spaces")
11271142
{

‎common/config/config.go‎

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ type Config struct {
3131
CityToCdnDomain map[string]string `env:"STARHUB_SERVER_CITY_TO_CDN_DOMAIN" default:""`
3232
UniqueServiceName string `env:"STARHUB_SERVER_UNIQUE_SERVICE_NAME" default:""`
3333
ServerFailureRedirectURL string `env:"STARHUB_SERVER_FAIL_REDIRECT_URL" default:"http://localhost:3000/errors/server-error"`
34+
NeedPhoneVerify bool `env:"STARHUB_SERVER_NEED_PHONE_VERIFY" default:"false"`
3435

3536
APIServer struct {
3637
Port int `env:"STARHUB_SERVER_SERVER_PORT" default:"8080"`
@@ -211,11 +212,12 @@ type Config struct {
211212
}
212213

213214
Accounting struct {
214-
Host string `env:"OPENCSG_ACCOUNTING_SERVER_HOST" default:"http://localhost"`
215-
Port int `env:"OPENCSG_ACCOUNTING_SERVER_PORT" default:"8086"`
216-
ChargingEnable bool `env:"OPENCSG_ACCOUNTING_CHARGING_ENABLE" default:"false"`
217-
SubscriptionCronExpression string `env:"OPENCSG_ACCOUNTING_SUBSCRIPTION_CRON_EXPRESSION" default:"*/5 * * * *"`
218-
ThresholdOfStopDeploy int `env:"OPENCSG_ACCOUNTING_THRESHOLD_OF_STOP_DEPLOY" default:"5000"`
215+
Host string `env:"OPENCSG_ACCOUNTING_SERVER_HOST" default:"http://localhost"`
216+
Port int `env:"OPENCSG_ACCOUNTING_SERVER_PORT" default:"8086"`
217+
ChargingEnable bool `env:"OPENCSG_ACCOUNTING_CHARGING_ENABLE" default:"false"`
218+
SubscriptionCronExpression string `env:"OPENCSG_ACCOUNTING_SUBSCRIPTION_CRON_EXPRESSION" default:"*/5 * * * *"`
219+
ExpiredPresentCronExpression string `env:"OPENCSG_ACCOUNTING_EXPIRED_PRESENT_CRON_EXPRESSION" default:"0 0 * * *"`
220+
ThresholdOfStopDeploy int `env:"OPENCSG_ACCOUNTING_THRESHOLD_OF_STOP_DEPLOY" default:"5000"`
219221
}
220222

221223
User struct {
@@ -334,12 +336,15 @@ type Config struct {
334336
IncreaseMultisyncRepoLimitCronExpression string `env:"STARHUB_SERVER_CRON_JOB_INCREASE_MULTISYNC_REPO_LIMIT_CRON_EXPRESSION" default:"0 0 * * *"`
335337
MigrateRepoPathCronExpression string `env:"STARHUB_SERVER_CRON_JOB_MIGRATE_REPO_PATH_CRON_EXPRESSION" default:"* 16-20 * * *"`
336338
DeletePendingDeletionCronExpression string `env:"STARHUB_SERVER_CRON_JOB_DELETE_PENDING_DELETION_CRON_EXPRESSION" default:"0 16-20 * * *"`
339+
ReleaseInvitationCreditCronExpression string `env:"STARHUB_SERVER_CRON_JOB_RELEASE_INVITATION_CREDIT_CRON_EXPRESSION" default:"0 0 5 * *"`
337340
}
338341

339342
Agent struct {
340-
AutoHubServiceHost string `env:"OPENCSG_AGENT_AUTOHUB_SERVICE_HOST" default:"http://internal.opencsg-stg.com:8190"`
341-
AgentHubServiceHost string `env:"OPENCSG_AGENT_AGENTHUB_SERVICE_HOST" default:""`
342-
AgentHubServiceToken string `env:"OPENCSG_AGENT_AGENTHUB_SERVICE_TOKEN" default:""`
343+
AutoHubServiceHost string `env:"OPENCSG_AGENT_AUTOHUB_SERVICE_HOST" default:"http://internal.opencsg-stg.com:8190"`
344+
AgentHubServiceHost string `env:"OPENCSG_AGENT_AGENTHUB_SERVICE_HOST" default:""`
345+
AgentHubServiceToken string `env:"OPENCSG_AGENT_AGENTHUB_SERVICE_TOKEN" default:""`
346+
CodeInstanceQuotaPerUser int `env:"STARHUB_SERVER_AGENT_CODE_INSTANCE_QUOTA_PER_USER" default:"5"`
347+
LangflowInstanceQuotaPerUser int `env:"STARHUB_SERVER_AGENT_LANGFLOW_INSTANCE_QUOTA_PER_USER" default:"5"`
343348
}
344349

345350
DataViewer struct {
@@ -419,6 +424,7 @@ type Config struct {
419424
DirectMailRegionId string `env:"STARHUB_SERVER_DIRECT_MAIL_REGION_ID" default:"cn-hangzhou"`
420425
MailerRechargeAdmin string `env:"STARHUB_SERVER_MAILER_RECHARGE_ADMIN" default:"[email protected]"`
421426
MailerWeeklyRechargesMail string `env:"STARHUB_SERVER_MAILER_WEEKLY_RECHARGES_MAIL" default:"[email protected]"`
427+
EmailInvoiceCreatedReceiver string `env:"STARHUB_SERVER_EMAIL_INVOICE_CREATED_RECEIVER" default:"[email protected]"`
422428
RepoSyncTimezone string `env:"STARHUB_SERVER_REPO_SYNC_TIMEZONE" default:"Asia/Shanghai"`
423429
NotificationRetryCount int `env:"STARHUB_SERVER_NOTIFIER_NOTIFICATION_RETRY_COUNT" default:"3"`
424430
BroadcastUserPageSize int `env:"STARHUB_SERVER_NOTIFIER_BROADCAST_USER_PAGE_SIZE" default:"100"`
@@ -494,7 +500,9 @@ type Config struct {
494500
HealthInterval int `env:"STARHUB_SERVER_LOGCOLLECTOR_HEALTH_INTERVAL" default:"5"`
495501
AcceptLabelPrefix string `env:"STARHUB_SERVER_LOGCOLLECTOR_ACCEPT_LABEL_PREFIX" default:"csghub_"`
496502
// the separator of log lines, default is "\\n" by client formats, "\n" sse auto newline
497-
LineSeparator string `env:"STARHUB_SERVER_LOGCOLLECTOR_LINE_SEPARATOR" default:"\\n"`
503+
LineSeparator string `env:"STARHUB_SERVER_LOGCOLLECTOR_LINE_SEPARATOR" default:"\\n"`
504+
MaxStoreTimeDay int `env:"STARHUB_SERVER_LOGCOLLECTOR_MAX_STORE_TIME_DAY" default:"7"`
505+
QueryLastReportTimeout int `env:"STARHUB_SERVER_LOGCOLLECTOR_QUERY_LAST_REPORT_TIMEOUT" default:"10"`
498506
}
499507

500508
Temporal struct {
@@ -503,6 +511,25 @@ type Config struct {
503511
MaxConcurrentWorkflowTaskExecutionSize int `env:"OPENCSG_TEMPORAL_MAX_CONCURRENT_WORKFLOW_TASK_EXECUTION_SIZE" default:"50"`
504512
}
505513

514+
APIRateLimiter struct {
515+
Enable bool `env:"STARHUB_SERVER_API_RATE_LIMITER_ENABLE" default:"false"`
516+
Limit int64 `env:"STARHUB_SERVER_API_RATE_LIMITER_LIMIT" default:"10"`
517+
Window int64 `env:"STARHUB_SERVER_API_RATE_LIMITER_WINDOW" default:"60"`
518+
}
519+
520+
APILocationCheck struct {
521+
Enable bool `env:"STARHUB_SERVER_API_LOCATION_CHECK_ENABLE" default:"false"`
522+
WhiteList []string `env:"STARHUB_SERVER_API_LOCATION_CHECK_WHITE_LIST" default:"[China,Hong Kong,Singapore]"`
523+
}
524+
525+
Captcha struct {
526+
ExceptionPaths []string `env:"STARHUB_SERVER_CAPTCHA_EXCEPTION_PATHS" default:"[/api/v1/broadcasts,/api/v1/notifications]"`
527+
}
528+
529+
GeoIP struct {
530+
DBFile string `env:"STARHUB_SERVER_GEOIP_DB_FILE" default:"/starhub-bin/GeoLite2-Country.mmdb"`
531+
}
532+
506533
Xnet struct {
507534
Endpoint string `env:"STARHUB_SERVER_XNET_ENDPOINT" default:"http://localhost:8097"`
508535
ApiKey string `env:"STARHUB_SERVER_XNET_API_KEY" default:"f3a7b9c1d6e5f8e2a1b5d4f9e6a2b8d7c3a4e2b1d9f6e7a8d2c5a7b4c1e3f5b8a1d4f9b7d6e2f8a5d3b1e7f9c6a8b2d1e4f7d5b6e9f2a4b3c8e1d7f995hd82hf"`

0 commit comments

Comments
 (0)