Add ratelimiter for fault injection handler

agent/handlers/task_server_setup.go

@@ -39,6 +39,8 @@ import (
 	tmdsv4 "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/handlers/v4"
 	"github.com/aws/amazon-ecs-agent/ecs-agent/utils/retry"
 	"github.com/cihub/seelog"
+	"github.com/didip/tollbooth"
+	"github.com/didip/tollbooth/limiter"
 	"github.com/gorilla/mux"
 )
 
@@ -93,7 +95,7 @@ func taskServerSetup(
 		taskProtectionClientFactory, metricsFactory)
 
 	// TODO: Future PR to pass in TMDS server router once all of the handlers have been implemented.
-	registerFaultHandlers(nil, tmdsAgentState, metricsFactory)
+	registerFaultHandlers(muxRouter, tmdsAgentState, metricsFactory)
 
 	return tmds.NewServer(auditLogger,
 		tmds.WithHandler(muxRouter),
@@ -205,51 +207,61 @@ func registerFaultHandlers(
 		return
 	}
 
+	rateLimiter := createRateLimiter()
+
 	// Setting up handler endpoints for network blackhole port fault injections
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.BlackHolePortFaultType),
-		handler.StartNetworkBlackholePort(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.StartNetworkBlackholePort()),
 	).Methods("PUT")
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.BlackHolePortFaultType),
-		handler.StopNetworkBlackHolePort(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.StopNetworkBlackHolePort()),
 	).Methods("DELETE")
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.BlackHolePortFaultType),
-		handler.CheckNetworkBlackHolePort(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.CheckNetworkBlackHolePort()),
 	).Methods("GET")
 
 	// Setting up handler endpoints for network latency fault injections
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.LatencyFaultType),
-		handler.StartNetworkLatency(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.StartNetworkLatency()),
 	).Methods("PUT")
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.LatencyFaultType),
-		handler.StopNetworkLatency(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.StopNetworkLatency()),
 	).Methods("DELETE")
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.LatencyFaultType),
-		handler.CheckNetworkLatency(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.CheckNetworkLatency()),
 	).Methods("GET")
 
 	// Setting up handler endpoints for network packet loss fault injections
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.PacketLossFaultType),
-		handler.StartNetworkPacketLoss(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.StartNetworkPacketLoss()),
 	).Methods("PUT")
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.PacketLossFaultType),
-		handler.StopNetworkPacketLoss(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.StopNetworkPacketLoss()),
 	).Methods("DELETE")
-	muxRouter.HandleFunc(
+	muxRouter.Handle(
 		fault.NetworkFaultPath(faulttype.PacketLossFaultType),
-		handler.CheckNetworkPacketLoss(),
+		tollbooth.LimitFuncHandler(rateLimiter, handler.CheckNetworkPacketLoss()),
 	).Methods("GET")
 
 	seelog.Debug("Successfully set up Fault TMDS handlers")
 }
 
+// Creates a tollbooth ratelimiter for the Fault Handler APIs
+func createRateLimiter() *limiter.Limiter {
+	lmt := tollbooth.NewLimiter(0.2, nil)
+	lmt.SetBurst(1)
+	lmt.SetMessage("You have reached maximum request limit")
+	return lmt
+}
+
 // ServeTaskHTTPEndpoint serves task/container metadata, task/container stats, IAM Role Credentials, and Agent APIs
 // for tasks being managed by the agent.
 func ServeTaskHTTPEndpoint(