Consume TMDS init function from ecs-agent module (#3663)

agent/config/config.go

@@ -44,9 +44,6 @@ const (
 	// AgentIntrospectionPort is used to serve the metadata about the agent and to query the tasks being managed by the agent.
 	AgentIntrospectionPort = 51678
 
-	// AgentCredentialsPort is used to serve the credentials for tasks.
-	AgentCredentialsPort = 51679
-
 	// AgentPrometheusExpositionPort is used to expose Prometheus metrics that can be scraped by a Prometheus server
 	AgentPrometheusExpositionPort = 51680
 

agent/config/config_unix.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/aws/amazon-ecs-agent/agent/dockerclient"
 	"github.com/aws/amazon-ecs-agent/agent/utils"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds"
 )
 
 const (
@@ -61,7 +62,7 @@ const (
 func DefaultConfig() Config {
 	return Config{
 		DockerEndpoint:                      "unix:///var/run/docker.sock",
-		ReservedPorts:                       []uint16{SSHPort, DockerReservedPort, DockerReservedSSLPort, AgentIntrospectionPort, AgentCredentialsPort},
+		ReservedPorts:                       []uint16{SSHPort, DockerReservedPort, DockerReservedSSLPort, AgentIntrospectionPort, tmds.Port},
 		ReservedPortsUDP:                    []uint16{},
 		DataDir:                             "/data/",
 		DataDirOnHost:                       "/var/lib/ecs",

agent/config/config_windows.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/aws/amazon-ecs-agent/agent/dockerclient"
 	"github.com/aws/amazon-ecs-agent/agent/utils"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds"
 
 	"github.com/cihub/seelog"
 	"github.com/hectane/go-acl/api"
@@ -103,7 +104,7 @@ func DefaultConfig() Config {
 			DockerReservedPort,
 			DockerReservedSSLPort,
 			AgentIntrospectionPort,
-			AgentCredentialsPort,
+			tmds.Port,
 			rdpPort,
 			rpcPort,
 			smbPort,

agent/config/config_windows_test.go

@@ -26,6 +26,7 @@ import (
 
 	"github.com/aws/amazon-ecs-agent/agent/dockerclient"
 	"github.com/aws/amazon-ecs-agent/agent/ec2"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds"
 
 	"github.com/hectane/go-acl/api"
 	"github.com/stretchr/testify/assert"
@@ -83,7 +84,7 @@ func TestConfigIAMTaskRolesReserves80(t *testing.T) {
 		DockerReservedPort,
 		DockerReservedSSLPort,
 		AgentIntrospectionPort,
-		AgentCredentialsPort,
+		tmds.Port,
 		rdpPort,
 		rpcPort,
 		smbPort,

agent/go.mod

@@ -11,7 +11,6 @@ require (
 	github.com/containernetworking/cni v0.8.1
 	github.com/containernetworking/plugins v0.9.1
 	github.com/deniswernert/udev v0.0.0-20170418162847-a12666f7b5a1
-	github.com/didip/tollbooth v4.0.2+incompatible
 	github.com/docker/docker v20.10.23+incompatible
 	github.com/docker/go-connections v0.4.0
 	github.com/docker/go-units v0.4.0
@@ -46,6 +45,7 @@ require (
 	github.com/containerd/continuity v0.3.0 // indirect
 	github.com/coreos/go-systemd/v22 v22.3.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/didip/tollbooth v4.0.2+incompatible // indirect
 	github.com/docker/distribution v2.8.1+incompatible // indirect
 	github.com/godbus/dbus/v5 v5.0.6 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect

agent/handlers/introspection_server_setup.go

@@ -27,6 +27,7 @@ import (
 	handlersutils "github.com/aws/amazon-ecs-agent/agent/handlers/utils"
 	v1 "github.com/aws/amazon-ecs-agent/agent/handlers/v1"
 	"github.com/aws/amazon-ecs-agent/agent/utils/retry"
+	logginghandler "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/logging"
 	"github.com/cihub/seelog"
 )
 
@@ -81,7 +82,7 @@ func introspectionServerSetup(containerInstanceArn *string, taskEngine handlersu
 
 	// Log all requests and then pass through to serverMux
 	loggingServeMux := http.NewServeMux()
-	loggingServeMux.Handle("/", LoggingHandler{serverMux})
+	loggingServeMux.Handle("/", logginghandler.NewLoggingHandler(serverMux))
 
 	wTimeout := writeTimeout
 	if cfg.EnableRuntimeStats.Enabled() {

agent/handlers/task_server_setup.go

@@ -16,15 +16,13 @@ package handlers
 import (
 	"context"
 	"net/http"
-	"strconv"
 	"time"
 
 	"github.com/aws/amazon-ecs-agent/agent/api"
 	"github.com/aws/amazon-ecs-agent/agent/config"
 	"github.com/aws/amazon-ecs-agent/agent/credentials"
 	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate"
 	agentAPITaskProtectionV1 "github.com/aws/amazon-ecs-agent/agent/handlers/agentapi/taskprotection/v1/handlers"
-	handlersutils "github.com/aws/amazon-ecs-agent/agent/handlers/utils"
 	v1 "github.com/aws/amazon-ecs-agent/agent/handlers/v1"
 	v2 "github.com/aws/amazon-ecs-agent/agent/handlers/v2"
 	v3 "github.com/aws/amazon-ecs-agent/agent/handlers/v3"
@@ -33,8 +31,8 @@ import (
 	"github.com/aws/amazon-ecs-agent/agent/stats"
 	"github.com/aws/amazon-ecs-agent/agent/utils/retry"
 	auditinterface "github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds"
 	"github.com/cihub/seelog"
-	"github.com/didip/tollbooth"
 	"github.com/gorilla/mux"
 )
 
@@ -61,7 +59,8 @@ func taskServerSetup(credentialsManager credentials.Manager,
 	vpcID string,
 	containerInstanceArn string,
 	apiEndpoint string,
-	acceptInsecureCert bool) *http.Server {
+	acceptInsecureCert bool) (*http.Server, error) {
+
 	muxRouter := mux.NewRouter()
 
 	// Set this to false so that for request like "//v3//metadata/task"
@@ -79,28 +78,13 @@ func taskServerSetup(credentialsManager credentials.Manager,
 
 	agentAPIV1HandlersSetup(muxRouter, state, credentialsManager, cluster, region, apiEndpoint, acceptInsecureCert)
 
-	limiter := tollbooth.NewLimiter(float64(steadyStateRate), nil)
-	limiter.SetOnLimitReached(handlersutils.LimitReachedHandler(auditLogger))
-	limiter.SetBurst(burstRate)
-
-	// Log all requests and then pass through to muxRouter.
-	loggingMuxRouter := mux.NewRouter()
-
-	// rootPath is a path for any traffic to this endpoint, "root" mux name will not be used.
-	rootPath := "/" + handlersutils.ConstructMuxVar("root", handlersutils.AnythingRegEx)
-	loggingMuxRouter.Handle(rootPath, tollbooth.LimitHandler(
-		limiter, NewLoggingHandler(muxRouter)))
-
-	loggingMuxRouter.SkipClean(false)
-
-	server := http.Server{
-		Addr:         "127.0.0.1:" + strconv.Itoa(config.AgentCredentialsPort),
-		Handler:      loggingMuxRouter,
-		ReadTimeout:  readTimeout,
-		WriteTimeout: writeTimeout,
-	}
-
-	return &server
+	return tmds.NewServer(auditLogger,
+		tmds.WithRouter(muxRouter),
+		tmds.WithListenAddress(tmds.AddressIPv4()),
+		tmds.WithReadTimeout(readTimeout),
+		tmds.WithWriteTimeout(writeTimeout),
+		tmds.WithSteadyStateRate(float64(steadyStateRate)),
+		tmds.WithBurstRate(burstRate))
 }
 
 // v2HandlersSetup adds all handlers in v2 package to the mux router.
@@ -200,9 +184,13 @@ func ServeTaskHTTPEndpoint(
 
 	auditLogger := audit.NewAuditLog(containerInstanceArn, cfg, logger)
 
-	server := taskServerSetup(credentialsManager, auditLogger, state, ecsClient, cfg.Cluster, cfg.AWSRegion, statsEngine,
+	server, err := taskServerSetup(credentialsManager, auditLogger, state, ecsClient, cfg.Cluster, cfg.AWSRegion, statsEngine,
 		cfg.TaskMetadataSteadyStateRate, cfg.TaskMetadataBurstRate, availabilityZone, vpcID, containerInstanceArn, cfg.APIEndpoint,
 		cfg.AcceptInsecureCert)
+	if err != nil {
+		seelog.Criticalf("Failed to set up Task Metadata Server: %v", err)
+		return
+	}
 
 	go func() {
 		<-ctx.Done()