Add TMDS initialization functionality to ecs-agent module (aws#3660)

Author: amogh09

Makefile

@@ -171,7 +171,7 @@ analyze-cover-profile-init: coverprofile-init.out
 	./scripts/analyze-cover-profile coverprofile-init.out
 
 run-integ-tests: test-registry gremlin container-health-check-image run-sudo-tests
-	ECS_LOGLEVEL=debug ${GOTEST} -tags integration -timeout=30m ./agent/...
+	ECS_LOGLEVEL=debug ${GOTEST} -tags integration -timeout=30m ./agent/... ./ecs-agent/...
 
 run-sudo-tests:
 	sudo -E ${GOTEST} -tags sudo -timeout=10m ./agent/...

ecs-agent/go.mod

@@ -4,17 +4,22 @@ go 1.19
 
 require (
 	github.com/aws/aws-sdk-go v1.36.0
+	github.com/cihub/seelog v0.0.0-20170130134532-f561c5e57575
+	github.com/didip/tollbooth v4.0.2+incompatible
 	github.com/golang/mock v1.4.1
+	github.com/gorilla/mux v1.8.0
 	github.com/stretchr/testify v1.7.0
 	golang.org/x/tools v0.6.0
 )
 
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	golang.org/x/mod v0.8.0 // indirect
 	golang.org/x/net v0.8.0 // indirect
 	golang.org/x/sys v0.6.0 // indirect
+	golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

ecs-agent/go.sum

@@ -1,14 +1,22 @@
 github.com/aws/aws-sdk-go v1.36.0 h1:CscTrS+szX5iu34zk2bZrChnGO/GMtUYgMK1Xzs2hYo=
 github.com/aws/aws-sdk-go v1.36.0/go.mod h1:hcU610XS61/+aQV88ixoOzUoG7v3b31pl2zKMmprdro=
+github.com/cihub/seelog v0.0.0-20170130134532-f561c5e57575 h1:kHaBemcxl8o/pQ5VM1c8PVE1PubbNx3mjUr09OqWGCs=
+github.com/cihub/seelog v0.0.0-20170130134532-f561c5e57575/go.mod h1:9d6lWj8KzO/fd/NrVaLscBKmPigpZpn5YawRPw+e3Yo=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/didip/tollbooth v4.0.2+incompatible h1:fVSa33JzSz0hoh2NxpwZtksAzAgd7zjmGO20HCZtF4M=
+github.com/didip/tollbooth v4.0.2+incompatible/go.mod h1:A9b0665CE6l1KmzpDws2++elm/CsuWBMa5Jv4WY0PEY=
 github.com/golang/mock v1.4.1 h1:ocYkMQY5RrXTYgXl7ICpV0IXwlEQGwKIsery4gyXa1U=
 github.com/golang/mock v1.4.1/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw=
+github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
+github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
 github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
 github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
@@ -33,6 +41,8 @@ golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e h1:EHBhcS0mlXEAVwNyO2dLfjToGsyY4j24pTs2ScHnX7s=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20190425150028-36563e24a262/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q=
 golang.org/x/tools v0.6.0 h1:BOw41kyTf3PuCW1pVQf8+Cyg8pMlkYB1oo9iJ6D/lKM=

ecs-agent/tmds/logging/logging_handler.go

@@ -0,0 +1,34 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package logging
+
+import (
+	"net/http"
+
+	"github.com/cihub/seelog"
+)
+
+// LoggingHandler is used to log all requests for an endpoint.
+type LoggingHandler struct{ h http.Handler }
+
+// NewLoggingHandler creates a new LoggingHandler object.
+func NewLoggingHandler(handler http.Handler) LoggingHandler {
+	return LoggingHandler{h: handler}
+}
+
+// ServeHTTP logs the method and remote address of the request.
+func (lh LoggingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	seelog.Debug("Handling http request", "method", r.Method, "from", r.RemoteAddr)
+	lh.h.ServeHTTP(w, r)
+}

ecs-agent/tmds/logging/logging_handler_test.go

@@ -0,0 +1,46 @@
+//go:build unit
+// +build unit
+
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+package logging
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type underlyingHandler struct{}
+
+func (h underlyingHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	fmt.Fprint(w, "Hello world")
+}
+
+// Tests that logging handler calls the underlying handler
+func TestLoggingHandler(t *testing.T) {
+	loggingHandler := LoggingHandler{underlyingHandler{}}
+
+	req, err := http.NewRequest("GET", "/", nil)
+	require.NoError(t, err)
+	res := httptest.NewRecorder()
+
+	loggingHandler.ServeHTTP(res, req)
+	assert.Equal(t, http.StatusOK, res.Code)
+	assert.Equal(t, "Hello world", res.Body.String())
+}

ecs-agent/tmds/server.go

@@ -0,0 +1,144 @@
+// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+package tmds
+
+import (
+	"errors"
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/logger/audit/request"
+	"github.com/aws/amazon-ecs-agent/ecs-agent/tmds/logging"
+	muxutils "github.com/aws/amazon-ecs-agent/ecs-agent/tmds/utils/mux"
+
+	"github.com/didip/tollbooth"
+	"github.com/gorilla/mux"
+)
+
+const (
+	// TMDS IP and port
+	IPv4 = "127.0.0.1"
+	Port = 51679
+)
+
+// IPv4 address for TMDS
+func AddressIPv4() string {
+	return fmt.Sprintf("%s:%d", IPv4, Port)
+}
+
+// Configuration for TMDS
+type Config struct {
+	listenAddress   string        // http server listen address
+	readTimeout     time.Duration // http server read timeout
+	writeTimeout    time.Duration // http server write timeout
+	steadyStateRate float64       // steady request rate limit
+	burstRate       int           // burst request rate limit
+	router          *mux.Router   // router with routes configured
+}
+
+// Function type for updating TMDS config
+type ConfigOpt func(*Config)
+
+// Set TMDS listen address
+func WithListenAddress(listenAddr string) ConfigOpt {
+	return func(c *Config) {
+		c.listenAddress = listenAddr
+	}
+}
+
+// Set TMDS read timeout
+func WithReadTimeout(readTimeout time.Duration) ConfigOpt {
+	return func(c *Config) {
+		c.readTimeout = readTimeout
+	}
+}
+
+// Set TMDS write timeout
+func WithWriteTimeout(writeTimeout time.Duration) ConfigOpt {
+	return func(c *Config) {
+		c.writeTimeout = writeTimeout
+	}
+}
+
+// Set TMDS steady request rate limit
+func WithSteadyStateRate(steadyStateRate float64) ConfigOpt {
+	return func(c *Config) {
+		c.steadyStateRate = steadyStateRate
+	}
+}
+
+// Set TMDS burst request rate limit
+func WithBurstRate(burstRate int) ConfigOpt {
+	return func(c *Config) {
+		c.burstRate = burstRate
+	}
+}
+
+// Set TMDS router
+func WithRouter(router *mux.Router) ConfigOpt {
+	return func(c *Config) {
+		c.router = router
+	}
+}
+
+// Create a new HTTP Task Metadata Server (TMDS)
+func NewServer(auditLogger audit.AuditLogger, options ...ConfigOpt) (*http.Server, error) {
+	config := new(Config)
+	for _, opt := range options {
+		opt(config)
+	}
+
+	return setup(auditLogger, config)
+}
+
+func setup(auditLogger audit.AuditLogger, config *Config) (*http.Server, error) {
+	if config.router == nil {
+		return nil, errors.New("router cannot be nil")
+	}
+
+	// Define a reqeuest rate limiter
+	limiter := tollbooth.
+		NewLimiter(config.steadyStateRate, nil).
+		SetOnLimitReached(limitReachedHandler(auditLogger)).
+		SetBurst(config.burstRate)
+
+	// Log all requests and then pass through to muxRouter.
+	loggingMuxRouter := mux.NewRouter()
+
+	// rootPath is a path for any traffic to this endpoint
+	rootPath := "/" + muxutils.ConstructMuxVar("root", muxutils.AnythingRegEx)
+	loggingMuxRouter.Handle(rootPath, tollbooth.LimitHandler(
+		limiter, logging.NewLoggingHandler(config.router)))
+
+	// explicitly enable path cleaning
+	loggingMuxRouter.SkipClean(false)
+
+	return &http.Server{
+		Addr:         config.listenAddress,
+		Handler:      loggingMuxRouter,
+		ReadTimeout:  config.readTimeout,
+		WriteTimeout: config.writeTimeout,
+	}, nil
+}
+
+// LimitReachedHandler logs the throttled request in the credentials audit log
+func limitReachedHandler(auditLogger audit.AuditLogger) func(http.ResponseWriter, *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		logRequest := request.LogRequest{
+			Request: r,
+		}
+		auditLogger.Log(logRequest, http.StatusTooManyRequests, "")
+	}
+}