Add TMDS initialization functionality to ecs-agent module #3660
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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()) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
singholt marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| ) | ||
|
|
||
| // 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) { | ||
|
There was a problem hiding this comment. You can see how this function will be consumed by agent module in the draft PR below. |
||
| 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. | ||
|
There was a problem hiding this comment. Server setup here is taken from agent module amazon-ecs-agent/agent/handlers/task_server_setup.go Lines 81 to 100 in 0202ad2 |
||
| 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) { | ||
|
There was a problem hiding this comment. This handler is taken from agent module amazon-ecs-agent/agent/handlers/utils/helpers.go Lines 128 to 135 in 0202ad2 |
||
| return func(w http.ResponseWriter, r *http.Request) { | ||
| logRequest := request.LogRequest{ | ||
| Request: r, | ||
| } | ||
| auditLogger.Log(logRequest, http.StatusTooManyRequests, "") | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This is taken from agent module.
amazon-ecs-agent/agent/handlers/logging_handler.go
Lines 22 to 34 in 0202ad2