Skip to content

Commit 5eeec98

Browse files
author
Thean Lim
committed
Migrate Websocket and TCS to v2 credentials
1 parent d1d87e2 commit 5eeec98

File tree

15 files changed

+106
-73
lines changed

15 files changed

+106
-73
lines changed

‎agent/app/agent.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -991,7 +991,7 @@ func (agent *ecsAgent) startAsyncRoutines(
991991
}
992992
go statsEngine.StartMetricsPublish()
993993

994-
session, err := reporter.NewDockerTelemetrySession(agent.containerInstanceARN, agent.credentialProvider, agent.cfg, deregisterInstanceEventStream,
994+
session, err := reporter.NewDockerTelemetrySession(agent.containerInstanceARN, agent.credentialsCache, agent.cfg, deregisterInstanceEventStream,
995995
client, taskEngine, telemetryMessages, healthMessages, doctor)
996996
if err != nil {
997997
seelog.Warnf("Error creating telemetry session: %v", err)
@@ -1102,7 +1102,7 @@ func (agent *ecsAgent) startACSSession(
11021102
acsSession := session.NewSession(agent.containerInstanceARN,
11031103
agent.cfg.Cluster,
11041104
client,
1105-
agent.credentialProvider,
1105+
agent.credentialsCache,
11061106
inactiveInstanceCB,
11071107
acsclient.NewACSClientFactory(),
11081108
metricsfactory.NewNopEntryFactory(),

‎agent/stats/reporter/reporter.go

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ import (
3030
tcshandler "github.com/aws/amazon-ecs-agent/ecs-agent/tcs/handler"
3131
"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
3232
"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
33-
"github.com/aws/aws-sdk-go/aws/credentials"
33+
34+
"github.com/aws/aws-sdk-go-v2/aws"
3435
)
3536

3637
const (
@@ -47,7 +48,7 @@ type DockerTelemetrySession struct {
4748
// tcshandler.TelemetrySession contains the logic to manage the TCSClient and corresponding websocket connection
4849
func NewDockerTelemetrySession(
4950
containerInstanceArn string,
50-
credentialProvider *credentials.Credentials,
51+
credentialsCache *aws.CredentialsCache,
5152
cfg *config.Config,
5253
deregisterInstanceEventStream *eventstream.EventStream,
5354
ecsClient ecs.ECSClient,
@@ -75,7 +76,7 @@ func NewDockerTelemetrySession(
7576
agentHash,
7677
containerRuntimeVersion,
7778
cfg.DisableMetrics.Enabled(),
78-
credentialProvider,
79+
credentialsCache,
7980
&wsclient.WSClientMinAgentConfig{
8081
AWSRegion: cfg.AWSRegion,
8182
AcceptInsecureCert: cfg.AcceptInsecureCert,

‎agent/stats/reporter/reporter_test.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,9 @@ import (
2626
"github.com/aws/amazon-ecs-agent/agent/version"
2727
"github.com/aws/amazon-ecs-agent/ecs-agent/doctor"
2828
"github.com/aws/amazon-ecs-agent/ecs-agent/eventstream"
29-
"github.com/aws/aws-sdk-go/aws/credentials"
29+
30+
"github.com/aws/aws-sdk-go-v2/aws"
31+
"github.com/aws/aws-sdk-go-v2/credentials"
3032
"github.com/golang/mock/gomock"
3133
"github.com/stretchr/testify/assert"
3234
)
@@ -41,7 +43,7 @@ const (
4143

4244
func TestNewDockerTelemetrySession(t *testing.T) {
4345
emptyDoctor, _ := doctor.NewDoctor([]doctor.Healthcheck{}, testCluster, testContainerInstanceArn)
44-
testCredentials := credentials.NewStaticCredentials("test-id", "test-secret", "test-token")
46+
testCredentials := credentials.NewStaticCredentialsProvider("test-id", "test-secret", "test-token")
4547
ctrl := gomock.NewController(t)
4648
defer ctrl.Finish()
4749
mockEngine := mock_engine.NewMockTaskEngine(ctrl)
@@ -94,7 +96,7 @@ func TestNewDockerTelemetrySession(t *testing.T) {
9496
t.Run(tc.name, func(t *testing.T) {
9597
dockerTelemetrySession, err := NewDockerTelemetrySession(
9698
testContainerInstanceArn,
97-
testCredentials,
99+
aws.NewCredentialsCache(testCredentials),
98100
tc.cfg,
99101
eventstream.NewEventStream("Deregister_Instance", context.Background()),
100102
nil,

‎ecs-agent/acs/client/acs_client.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ import (
2626
"github.com/aws/amazon-ecs-agent/ecs-agent/logger"
2727
"github.com/aws/amazon-ecs-agent/ecs-agent/metrics"
2828
"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
29-
"github.com/aws/aws-sdk-go/aws/credentials"
29+
"github.com/aws/aws-sdk-go-v2/aws"
3030
)
3131

3232
// clientServer implements ClientServer for acs.
@@ -46,14 +46,14 @@ func NewACSClientFactory() wsclient.ClientFactory {
4646
// The returned struct should have both 'Connect' and 'Serve' called upon it
4747
// before being used.
4848
func (*acsClientFactory) New(url string,
49-
credentialProvider *credentials.Credentials,
49+
credentialsCache *aws.CredentialsCache,
5050
rwTimeout time.Duration,
5151
cfg *wsclient.WSClientMinAgentConfig,
5252
metricsFactory metrics.EntryFactory) wsclient.ClientServer {
5353

5454
cs := &clientServer{}
5555
cs.URL = url
56-
cs.CredentialProvider = credentialProvider
56+
cs.CredentialsCache = credentialsCache
5757
cs.Cfg = cfg
5858
cs.ServiceError = &acsError{}
5959
cs.RequestHandlers = make(map[string]wsclient.RequestHandler)

‎ecs-agent/acs/client/acs_client_test.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,9 @@ import (
3232
mock_wsconn "github.com/aws/amazon-ecs-agent/ecs-agent/wsclient/wsconn/mock"
3333

3434
"github.com/aws/aws-sdk-go-v2/aws"
35+
"github.com/aws/aws-sdk-go-v2/credentials"
3536
"github.com/aws/aws-sdk-go-v2/service/acs"
3637
acstypes "github.com/aws/aws-sdk-go-v2/service/acs/types"
37-
"github.com/aws/aws-sdk-go/aws/credentials"
3838
"github.com/golang/mock/gomock"
3939
"github.com/gorilla/websocket"
4040
"github.com/stretchr/testify/assert"
@@ -110,7 +110,7 @@ const (
110110
rwTimeout = time.Second
111111
)
112112

113-
var testCreds = credentials.NewStaticCredentials("test-id", "test-secret", "test-token")
113+
var testCreds = credentials.NewStaticCredentialsProvider("test-id", "test-secret", "test-token")
114114

115115
var testCfg = &wsclient.WSClientMinAgentConfig{
116116
AcceptInsecureCert: true,
@@ -267,7 +267,7 @@ func TestConnect(t *testing.T) {
267267
t.Fatal(<-serverErr)
268268
}()
269269

270-
cs := testACSClientFactory.New(server.URL, testCreds, rwTimeout, testCfg, metrics.NewNopEntryFactory())
270+
cs := testACSClientFactory.New(server.URL, aws.NewCredentialsCache(testCreds), rwTimeout, testCfg, metrics.NewNopEntryFactory())
271271
// Wait for up to a second for the mock server to launch
272272
for i := 0; i < 100; i++ {
273273
_, err = cs.Connect(metrics.ACSDisconnectTimeoutMetricName, wsclient.DisconnectTimeout, wsclient.DisconnectJitterMax)
@@ -338,15 +338,15 @@ func TestConnectClientError(t *testing.T) {
338338
}))
339339
defer testServer.Close()
340340

341-
cs := testACSClientFactory.New(testServer.URL, testCreds, rwTimeout, testCfg, metrics.NewNopEntryFactory())
341+
cs := testACSClientFactory.New(testServer.URL, aws.NewCredentialsCache(testCreds), rwTimeout, testCfg, metrics.NewNopEntryFactory())
342342
_, err := cs.Connect(metrics.ACSDisconnectTimeoutMetricName, wsclient.DisconnectTimeout, wsclient.DisconnectJitterMax)
343343
_, ok := err.(*wsclient.WSError)
344344
assert.True(t, ok, "Connect error expected to be a WSError type")
345345
assert.EqualError(t, err, "InvalidClusterException: Invalid cluster")
346346
}
347347

348348
func testCS(conn *mock_wsconn.MockWebsocketConn) wsclient.ClientServer {
349-
foo := testACSClientFactory.New("localhost:443", testCreds, rwTimeout, testCfg, metrics.NewNopEntryFactory())
349+
foo := testACSClientFactory.New("localhost:443", aws.NewCredentialsCache(testCreds), rwTimeout, testCfg, metrics.NewNopEntryFactory())
350350
cs := foo.(*clientServer)
351351
cs.SetConnection(conn)
352352
return cs

‎ecs-agent/acs/session/session.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,8 @@ import (
3232
"github.com/aws/amazon-ecs-agent/ecs-agent/utils/retry"
3333
"github.com/aws/amazon-ecs-agent/ecs-agent/utils/ttime"
3434
"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
35-
"github.com/aws/aws-sdk-go/aws/credentials"
35+
36+
"github.com/aws/aws-sdk-go-v2/aws"
3637
)
3738

3839
const (
@@ -71,7 +72,7 @@ type Session interface {
7172
type session struct {
7273
containerInstanceARN string
7374
cluster string
74-
credentialsProvider *credentials.Credentials
75+
credentialsCache *aws.CredentialsCache
7576
ecsClient ecs.ECSClient
7677
inactiveInstanceCB func()
7778
agentVersion string
@@ -106,7 +107,7 @@ type session struct {
106107
func NewSession(containerInstanceARN string,
107108
cluster string,
108109
ecsClient ecs.ECSClient,
109-
credentialsProvider *credentials.Credentials,
110+
credentialsCache *aws.CredentialsCache,
110111
inactiveInstanceCB func(),
111112
clientFactory wsclient.ClientFactory,
112113
metricsFactory metrics.EntryFactory,
@@ -132,7 +133,7 @@ func NewSession(containerInstanceARN string,
132133
containerInstanceARN: containerInstanceARN,
133134
cluster: cluster,
134135
ecsClient: ecsClient,
135-
credentialsProvider: credentialsProvider,
136+
credentialsCache: credentialsCache,
136137
inactiveInstanceCB: inactiveInstanceCB,
137138
clientFactory: clientFactory,
138139
metricsFactory: metricsFactory,
@@ -247,7 +248,7 @@ func (s *session) startSessionOnce(ctx context.Context) error {
247248

248249
client := s.clientFactory.New(
249250
s.acsURL(acsEndpoint),
250-
s.credentialsProvider,
251+
s.credentialsCache,
251252
wsRWTimeout,
252253
s.minAgentConfig,
253254
s.metricsFactory)

‎ecs-agent/acs/session/session_test.go

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,9 @@ import (
4444
mock_retry "github.com/aws/amazon-ecs-agent/ecs-agent/utils/retry/mock"
4545
"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
4646
mock_wsclient "github.com/aws/amazon-ecs-agent/ecs-agent/wsclient/mock"
47-
"github.com/aws/aws-sdk-go/aws/credentials"
47+
48+
"github.com/aws/aws-sdk-go-v2/aws"
49+
"github.com/aws/aws-sdk-go-v2/credentials"
4850
"github.com/golang/mock/gomock"
4951
"github.com/gorilla/websocket"
5052
"github.com/pkg/errors"
@@ -180,7 +182,7 @@ const (
180182

181183
var inactiveInstanceError = errors.New("InactiveInstanceException")
182184
var noopFunc = func() {}
183-
var testCreds = credentials.NewStaticCredentials("test-id", "test-secret", "test-token")
185+
var testCreds = credentials.NewStaticCredentialsProvider("test-id", "test-secret", "test-token")
184186
var testMinAgentConfig = &wsclient.WSClientMinAgentConfig{
185187
AcceptInsecureCert: true,
186188
AWSRegion: "us-west-2",
@@ -973,7 +975,7 @@ func TestSessionDoesntLeakGoroutines(t *testing.T) {
973975
go func() {
974976
acsSession := session{
975977
containerInstanceARN: testconst.ContainerInstanceARN,
976-
credentialsProvider: testCreds,
978+
credentialsCache: aws.NewCredentialsCache(testCreds),
977979
dockerVersion: dockerVersion,
978980
minAgentConfig: testMinAgentConfig,
979981
ecsClient: ecsClient,
@@ -1052,7 +1054,7 @@ func TestStartSessionHandlesRefreshCredentialsMessages(t *testing.T) {
10521054
acsSession := NewSession(testconst.ContainerInstanceARN,
10531055
testconst.ClusterARN,
10541056
ecsClient,
1055-
testCreds,
1057+
aws.NewCredentialsCache(testCreds),
10561058
noopFunc,
10571059
acsclient.NewACSClientFactory(),
10581060
metricsfactory.NewNopEntryFactory(),
@@ -1318,7 +1320,7 @@ func TestStartSessionHandlesAttachResourceMessages(t *testing.T) {
13181320
acsSession := NewSession(testconst.ContainerInstanceARN,
13191321
testconst.ClusterARN,
13201322
ecsClient,
1321-
testCreds,
1323+
aws.NewCredentialsCache(testCreds),
13221324
noopFunc,
13231325
acsclient.NewACSClientFactory(),
13241326
metricsfactory.NewNopEntryFactory(),

‎ecs-agent/tcs/client/client.go

Lines changed: 10 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -28,10 +28,9 @@ import (
2828
"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
2929
"github.com/aws/amazon-ecs-agent/ecs-agent/utils"
3030
"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
31-
"github.com/aws/aws-sdk-go/aws"
32-
"github.com/aws/aws-sdk-go/aws/credentials"
33-
"github.com/aws/aws-sdk-go/private/protocol/json/jsonutil"
3431

32+
"github.com/aws/aws-sdk-go-v2/aws"
33+
"github.com/aws/aws-sdk-go/private/protocol/json/jsonutil"
3534
"github.com/cihub/seelog"
3635
"github.com/pborman/uuid"
3736
)
@@ -72,7 +71,7 @@ func New(url string,
7271
doctor *doctor.Doctor,
7372
disableResourceMetrics bool,
7473
publishMetricsInterval time.Duration,
75-
credentialProvider *credentials.Credentials,
74+
credentialsCache *aws.CredentialsCache,
7675
rwTimeout time.Duration,
7776
metricsMessages <-chan ecstcs.TelemetryMessage,
7877
healthMessages <-chan ecstcs.HealthMessage,
@@ -88,9 +87,9 @@ func New(url string,
8887
ClientServerImpl: wsclient.ClientServerImpl{
8988
URL: url,
9089
Cfg: cfg,
91-
CredentialProvider: credentialProvider,
90+
CredentialsCache: credentialsCache,
9291
RWTimeout: rwTimeout,
93-
MakeRequestHook: signRequestFunc(url, cfg.AWSRegion, credentialProvider),
92+
MakeRequestHook: signRequestFunc(url, cfg.AWSRegion, credentialsCache),
9493
TypeDecoder: NewTCSDecoder(),
9594
RequestHandlers: make(map[string]wsclient.RequestHandler),
9695
MetricsFactory: metricsFactory,
@@ -368,10 +367,10 @@ func copyServiceConnectMetrics(scMetrics []*ecstcs.GeneralMetricsWrapper) []*ecs
368367
// copyHealthMetadata performs a deep copy of HealthMetadata object
369368
func copyHealthMetadata(metadata *ecstcs.HealthMetadata, fin bool) *ecstcs.HealthMetadata {
370369
return &ecstcs.HealthMetadata{
371-
Cluster: aws.String(aws.StringValue(metadata.Cluster)),
372-
ContainerInstance: aws.String(aws.StringValue(metadata.ContainerInstance)),
370+
Cluster: aws.String(aws.ToString(metadata.Cluster)),
371+
ContainerInstance: aws.String(aws.ToString(metadata.ContainerInstance)),
373372
Fin: aws.Bool(fin),
374-
MessageId: aws.String(aws.StringValue(metadata.MessageId)),
373+
MessageId: aws.String(aws.ToString(metadata.MessageId)),
375374
}
376375
}
377376

@@ -492,7 +491,7 @@ func (cs *tcsClientServer) Close() error {
492491
}
493492

494493
// signRequestFunc is a MakeRequestHookFunc that signs each generated request
495-
func signRequestFunc(url, region string, credentialProvider *credentials.Credentials) wsclient.MakeRequestHookFunc {
494+
func signRequestFunc(url, region string, credentialsCache *aws.CredentialsCache) wsclient.MakeRequestHookFunc {
496495
return func(payload []byte) ([]byte, error) {
497496
reqBody := bytes.NewReader(payload)
498497

@@ -501,7 +500,7 @@ func signRequestFunc(url, region string, credentialProvider *credentials.Credent
501500
return nil, err
502501
}
503502

504-
err = utils.SignHTTPRequest(request, region, "ecs", credentialProvider, reqBody)
503+
err = utils.SignHTTPRequest(request, region, "ecs", credentialsCache, reqBody)
505504
if err != nil {
506505
return nil, err
507506
}

‎ecs-agent/tcs/client/client_test.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,9 @@ import (
3535
"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
3636
"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
3737
mock_wsconn "github.com/aws/amazon-ecs-agent/ecs-agent/wsclient/wsconn/mock"
38-
"github.com/aws/aws-sdk-go/aws"
39-
"github.com/aws/aws-sdk-go/aws/credentials"
38+
39+
"github.com/aws/aws-sdk-go-v2/aws"
40+
"github.com/aws/aws-sdk-go-v2/credentials"
4041
"github.com/golang/mock/gomock"
4142
"github.com/stretchr/testify/assert"
4243
)
@@ -99,7 +100,7 @@ func (fc *falseHealthcheck) GetLastHealthcheckTime() time.Time {
99100
return time.Date(1974, time.May, 19, 1, 2, 3, 4, time.UTC)
100101
}
101102

102-
var testCreds = credentials.NewStaticCredentials("test-id", "test-secret", "test-token")
103+
var testCreds = credentials.NewStaticCredentialsProvider("test-id", "test-secret", "test-token")
103104

104105
var emptyDoctor, _ = doctor.NewDoctor([]doctor.Healthcheck{}, "test-cluster", "this:is:an:instance:arn")
105106

@@ -647,7 +648,7 @@ func testCS(conn *mock_wsconn.MockWebsocketConn, metricsMessages <-chan ecstcs.T
647648
AcceptInsecureCert: true,
648649
}
649650
cs := New("https://aws.amazon.com/ecs", cfg, emptyDoctor, false, testPublishMetricsInterval,
650-
testCreds, rwTimeout, metricsMessages, healthMessages, metrics.NewNopEntryFactory()).(*tcsClientServer)
651+
aws.NewCredentialsCache(testCreds), rwTimeout, metricsMessages, healthMessages, metrics.NewNopEntryFactory()).(*tcsClientServer)
651652
cs.SetConnection(conn)
652653
return cs
653654
}
@@ -718,7 +719,7 @@ func TestHealthToPublishHealthRequests(t *testing.T) {
718719
IsDocker: true,
719720
}
720721

721-
cs := New("", cfg, emptyDoctor, true, testPublishMetricsInterval, testCreds, rwTimeout, nil, nil, metrics.NewNopEntryFactory())
722+
cs := New("", cfg, emptyDoctor, true, testPublishMetricsInterval, aws.NewCredentialsCache(testCreds), rwTimeout, nil, nil, metrics.NewNopEntryFactory())
722723
cs.SetConnection(conn)
723724

724725
testMetadata := &ecstcs.HealthMetadata{

‎ecs-agent/tcs/handler/handler.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,8 @@ import (
3030
"github.com/aws/amazon-ecs-agent/ecs-agent/tcs/model/ecstcs"
3131
"github.com/aws/amazon-ecs-agent/ecs-agent/utils/retry"
3232
"github.com/aws/amazon-ecs-agent/ecs-agent/wsclient"
33-
"github.com/aws/aws-sdk-go/aws/credentials"
33+
34+
"github.com/aws/aws-sdk-go-v2/aws"
3435
"github.com/cihub/seelog"
3536
)
3637

@@ -61,7 +62,7 @@ type telemetrySession struct {
6162
agentHash string
6263
containerRuntimeVersion string
6364
disableMetrics bool
64-
credentialsProvider *credentials.Credentials
65+
credentialsCache *aws.CredentialsCache
6566
cfg *wsclient.WSClientMinAgentConfig
6667
deregisterInstanceEventStream *eventstream.EventStream
6768
heartbeatTimeout time.Duration
@@ -82,7 +83,7 @@ func NewTelemetrySession(
8283
agentHash string,
8384
containerRuntimeVersion string,
8485
disableMetrics bool,
85-
credentialsProvider *credentials.Credentials,
86+
credentialsCache *aws.CredentialsCache,
8687
cfg *wsclient.WSClientMinAgentConfig,
8788
deregisterInstanceEventStream *eventstream.EventStream,
8889
heartbeatTimeout time.Duration,
@@ -102,7 +103,7 @@ func NewTelemetrySession(
102103
agentHash: agentHash,
103104
containerRuntimeVersion: containerRuntimeVersion,
104105
disableMetrics: disableMetrics,
105-
credentialsProvider: credentialsProvider,
106+
credentialsCache: credentialsCache,
106107
cfg: cfg,
107108
deregisterInstanceEventStream: deregisterInstanceEventStream,
108109
metricsChannel: metricsChannel,
@@ -158,7 +159,7 @@ func (session *telemetrySession) StartTelemetrySession(ctx context.Context) erro
158159
tcsEndpointUrl := formatURL(endpoint, session.cluster, session.containerInstanceArn, session.agentVersion,
159160
session.agentHash, containerRuntime, session.containerRuntimeVersion)
160161
client := tcsclient.New(tcsEndpointUrl, session.cfg, session.doctor, session.disableMetrics, tcsclient.DefaultContainerMetricsPublishInterval,
161-
session.credentialsProvider, wsRWTimeout, session.metricsChannel, session.healthChannel, session.metricsFactory)
162+
session.credentialsCache, wsRWTimeout, session.metricsChannel, session.healthChannel, session.metricsFactory)
162163
defer client.Close()
163164

164165
if session.deregisterInstanceEventStream != nil {

0 commit comments

Comments
 (0)