@@ -629,7 +629,7 @@ func (a *App) getContractBills(r *http.Request) (interface{}, mw.Response) {
629629// @license.name Apache 2.0
630630// @license.url http://www.apache.org/licenses/LICENSE-2.0.html
631631// @BasePath /
632- func Setup (router * mux.Router , gitCommit string , cl DBClient , relayClient rmb.Client , idxIntervals map [string ]uint ) error {
632+ func Setup (router * mux.Router , gitCommit string , cl DBClient , relayClient rmb.Client , idxIntervals map [string ]uint , rateLimitRPS int ) error {
633633
634634 a := App {
635635 cl : cl ,
@@ -638,32 +638,41 @@ func Setup(router *mux.Router, gitCommit string, cl DBClient, relayClient rmb.Cl
638638 idxIntervals : idxIntervals ,
639639 }
640640
641- router .HandleFunc ("/farms" , mw .AsHandlerFunc (a .listFarms ))
642- router .HandleFunc ("/stats" , mw .AsHandlerFunc (a .getStats ))
641+ // Create rate limiter middleware if rate limiting is enabled
642+ var rateLimiter * mw.RateLimiterMiddleware
643+ if rateLimitRPS > 0 {
644+ rateLimiter = mw .NewRateLimiterMiddleware (rateLimitRPS )
645+ log .Info ().Int ("rate_limit_rps" , rateLimitRPS ).Msg ("Rate limiting enabled" )
646+ } else {
647+ log .Info ().Msg ("Rate limiting disabled" )
648+ }
649+
650+ router .HandleFunc ("/farms" , mw .WithRateLimit (rateLimiter , a .listFarms ))
651+ router .HandleFunc ("/stats" , mw .WithRateLimit (rateLimiter , a .getStats ))
643652
644- router .HandleFunc ("/twins" , mw .AsHandlerFunc ( a .listTwins ))
645- router .HandleFunc ("/twins/{twin_id:[0-9]+}/consumption" , mw .AsHandlerFunc ( a .getTwinConsumption ))
653+ router .HandleFunc ("/twins" , mw .WithRateLimit ( rateLimiter , a .listTwins ))
654+ router .HandleFunc ("/twins/{twin_id:[0-9]+}/consumption" , mw .WithRateLimit ( rateLimiter , a .getTwinConsumption ))
646655
647- router .HandleFunc ("/nodes" , mw .AsHandlerFunc ( a .getNodes ))
648- router .HandleFunc ("/nodes/{node_id:[0-9]+}" , mw .AsHandlerFunc ( a .getNode ))
649- router .HandleFunc ("/nodes/{node_id:[0-9]+}/status" , mw .AsHandlerFunc ( a .getNodeStatus ))
650- router .HandleFunc ("/nodes/{node_id:[0-9]+}/statistics" , mw .AsHandlerFunc ( a .getNodeStatistics ))
651- router .HandleFunc ("/nodes/{node_id:[0-9]+}/gpu" , mw .AsHandlerFunc ( a .getNodeGpus ))
656+ router .HandleFunc ("/nodes" , mw .WithRateLimit ( rateLimiter , a .getNodes ))
657+ router .HandleFunc ("/nodes/{node_id:[0-9]+}" , mw .WithRateLimit ( rateLimiter , a .getNode ))
658+ router .HandleFunc ("/nodes/{node_id:[0-9]+}/status" , mw .WithRateLimit ( rateLimiter , a .getNodeStatus ))
659+ router .HandleFunc ("/nodes/{node_id:[0-9]+}/statistics" , mw .WithRateLimit ( rateLimiter , a .getNodeStatistics ))
660+ router .HandleFunc ("/nodes/{node_id:[0-9]+}/gpu" , mw .WithRateLimit ( rateLimiter , a .getNodeGpus ))
652661
653- router .HandleFunc ("/gateways" , mw .AsHandlerFunc ( a .getGateways ))
654- router .HandleFunc ("/gateways/{node_id:[0-9]+}" , mw .AsHandlerFunc ( a .getGateway ))
655- router .HandleFunc ("/gateways/{node_id:[0-9]+}/status" , mw .AsHandlerFunc ( a .getNodeStatus ))
662+ router .HandleFunc ("/gateways" , mw .WithRateLimit ( rateLimiter , a .getGateways ))
663+ router .HandleFunc ("/gateways/{node_id:[0-9]+}" , mw .WithRateLimit ( rateLimiter , a .getGateway ))
664+ router .HandleFunc ("/gateways/{node_id:[0-9]+}/status" , mw .WithRateLimit ( rateLimiter , a .getNodeStatus ))
656665
657- router .HandleFunc ("/contracts" , mw .AsHandlerFunc ( a .listContracts ))
658- router .HandleFunc ("/contracts/{contract_id:[0-9]+}" , mw .AsHandlerFunc ( a .getContract ))
659- router .HandleFunc ("/contracts/{contract_id:[0-9]+}/bills" , mw .AsHandlerFunc ( a .getContractBills ))
666+ router .HandleFunc ("/contracts" , mw .WithRateLimit ( rateLimiter , a .listContracts ))
667+ router .HandleFunc ("/contracts/{contract_id:[0-9]+}" , mw .WithRateLimit ( rateLimiter , a .getContract ))
668+ router .HandleFunc ("/contracts/{contract_id:[0-9]+}/bills" , mw .WithRateLimit ( rateLimiter , a .getContractBills ))
660669
661- router .HandleFunc ("/public_ips" , mw .AsHandlerFunc ( a .GetPublicIps ))
670+ router .HandleFunc ("/public_ips" , mw .WithRateLimit ( rateLimiter , a .GetPublicIps ))
662671
663- router .HandleFunc ("/" , mw .AsHandlerFunc ( a .indexPage (router )))
664- router .HandleFunc ("/ping" , mw .AsHandlerFunc ( a .ping ))
665- router .HandleFunc ("/version" , mw .AsHandlerFunc ( a .version ))
666- router .HandleFunc ("/health" , mw .AsHandlerFunc ( a .health ))
672+ router .HandleFunc ("/" , mw .WithRateLimit ( rateLimiter , a .indexPage (router )))
673+ router .HandleFunc ("/ping" , mw .WithRateLimit ( rateLimiter , a .ping ))
674+ router .HandleFunc ("/version" , mw .WithRateLimit ( rateLimiter , a .version ))
675+ router .HandleFunc ("/health" , mw .WithRateLimit ( rateLimiter , a .health ))
667676 router .PathPrefix ("/swagger/" ).Handler (httpSwagger .WrapHandler )
668677
669678 return nil
0 commit comments